[백준/Java] 2143 두 배열의 합

https://www.acmicpc.net/problem/2143

문제 요약

길이 n인 배열 A, 길이 m인 배열 B에 대해, 각 배열의 부분합의 총합이 T가 되는 모든 경우의 개수를 얻고자 한다.
T는 최소 -10억, 최대 10억이다. n과 m은 최대 1000이다.
배열의 각 원소는 절대값 100만 이하의 정수이다.

---

문제 접근

이 문제는 HashMap을 이용하여 문제를 해결할 수 있다.

먼저 첫번째 배열에 대해 모든 부분합의 빈도 수를 기록한다.
그 다음 두번째 배열의 모든 부분합의 빈도 수를 기록한다.
빈도수를 이용해 두 부분합을 더해서 총합이 T가 되는 경우를 구하면 된다.
이는 A의 부분합 빈도수를 훑으면서, 그 부분합과 T의 차가 B의 부분합에 존재하는지 확인하여 수행할 수 있다.
만약 존재한다면 빈도수 * 빈도수를 통해 모든 경우를 얻을 수 있다.

다른 방법으로는 투포인터나 이분탐색을 사용할 수도 있다.
투포인터로 문제를 해결하는 방법은 두 배열의 부배열 합을 각각 모두 구하고, 정렬한 뒤 각각 앞 뒤로 투포인터를 사용하는 것이다.
이분탐색으로 문제를 해결하는 방법은 똑같이 두 부배열의 합을 모두 구하고, 정렬한 뒤 A의 카운트, B의 카운트를 구해서 계산하는 것이다.

어떻게 문제에 접근하든, 일단 두 배열의 부배열의 합을 각각 모두 구해야 한다는 점은 동일하다.

---

풀이

기본적인 아이디어는 다음과 같다.

1. (정리 예정)

이를 구현한 코드는 다음과 같다.

import java.util.*;
import java.io.*;

public class Main {
    static int T, n, m;
    static int[] A, B;
    static HashMap<Long, Long> sumFreqOfA, sumFreqOfB;
    static long answer = 0;
    
    // 입력 받기
    public static void getInput() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        
        T = Integer.parseInt(br.readLine());
        n = Integer.parseInt(br.readLine());
        
        A = new int[n];
        StringTokenizer stA = new StringTokenizer(br.readLine());
        for (int i = 0; i < n; i++) {
            A[i] = Integer.parseInt(stA.nextToken());
        }
        
        m = Integer.parseInt(br.readLine());
        
        B = new int[m];
        StringTokenizer stB = new StringTokenizer(br.readLine());
        for (int j = 0; j < m; j++) {
            B[j] = Integer.parseInt(stB.nextToken());
        }
    }
    
    // 부분합 구하는 함수
    public static long calPartSum(int start, int end, int[] arr) {
    	long sum = 0;
    	
    	for (int i = start; i < end; i++) {
    		sum += arr[i];
    	}
    	
    	return sum;
    }
    
    // 빈도수 해쉬맵 계산
    public static void getSumFreq() {
    	sumFreqOfA = new HashMap<Long, Long>();
    	sumFreqOfB = new HashMap<Long, Long>();
    	
    	// A의 부분합 빈도수 구하기
    	// i는 부분합의 길이, j는 부분합 계산의 시작 인덱스
    	for (int i = 1; i <= n; i++) {
    		for (int j = 0; i+j <= n; j++) {
    			long partSum = calPartSum(j, i+j, A);
    			if (sumFreqOfA.containsKey(partSum)) {
    				Long former = sumFreqOfA.get(partSum);
    				sumFreqOfA.replace(partSum, former+1);
    			} else {
    				sumFreqOfA.put(partSum, 1L);
    			}
    		}
    	}
    	
    	// B의 빈도수 구하기
    	for (int x = 1; x <= m; x++) {
    		for (int y = 0; x+y <= m; y++) {
    			long partSum = calPartSum(y, x+y, B);
    			if (sumFreqOfB.containsKey(partSum)) {
    				Long former = sumFreqOfB.get(partSum);
    				sumFreqOfB.replace(partSum, former+1);
    			} else {
    				sumFreqOfB.put(partSum, 1L);
    			}
    		}
    	}
    }
    
    // 빈도수 해쉬맵으로 합이 T가 되는 경우 계산
    public static long calAnswer() {
    	long count = 0;
    	
    	for (Map.Entry<Long, Long> element : sumFreqOfA.entrySet()) {
    		long nowPartSumOfA = element.getKey();
    		long nowFreqSumOfA = element.getValue();
    		
    		long diff = T - nowPartSumOfA;
    		
    		long freqSumOfB = sumFreqOfB.getOrDefault(diff, 0L);
    		
    		count += nowFreqSumOfA * freqSumOfB;
    	}
    	
    	return count;
    }
    
    
    public static void main(String[] args) throws IOException {
        getInput();
        
        getSumFreq();
        
        answer = calAnswer();
        
        System.out.println(answer);
    }
}

---

회고

틀렸던 이유
합이 int를 넘어갈 건 예상했는데, 카운트가 넘어가는 것은 예상 못했다.
생각해보니 카운트를 부분합 빈도의 곱으로 구하니 최대 n^2 x m^2 이라서 100만 x 100만이 될 수 있으므로 int 범위를 넘어갈 수 있다.