백준 2143번: 두 배열의 합

문제 설명

• https://www.acmicpc.net/problem/2143
• A배열의 연속된 구간과 B배열의 연속된 구간을 합쳤을 때 T가 나오게 되는 경우의 수를 구하는 문제입니다.

접근법

• A배열의 부분합으로 가능한 경우의 수와 B배열의 부분합으로 가능한 경우의 수를 모두 구한 다음 둘을 비교합니다.
  • 가능한 부분합을 모두 구하는 시간복잡도는 각각 O(n^2)과 O(m^2)입니다.
  • 이때 구해진 부분합을 모두 비교하면 O(1000^4)으로 시간초과가 발생합니다.
    값 자체를 배열이나 리스트로 가지고 있는것보다, 어떤 값이 몇번 나왔는지를 나타내는 해시를 활용하면 시간복잡도를 줄일 수 있습니다.
    • 아래 코드를 보시면 알겠지만 저는 이러한 방법으로 문제를 통과했지만, 이 또한 모든 부분합이 한번만 등장하는 최악의 경우에는 시간초과가 발생할 거 같습니다.
    • 좀 더 올바른 풀이는 contains대신 이분탐색을 활용해 시간복잡도를 줄여야 할 거 같습니다.

정답

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

public class Main {
	static int answer = 0;
	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int T = Integer.parseInt(br.readLine());
		int n = Integer.parseInt(br.readLine());

		// A누적합
		int[] cumSumA = new int[n+1];
		StringTokenizer st = new StringTokenizer(br.readLine());
		for (int i = 1; i <= n; i++) {
			cumSumA[i] = Integer.parseInt(st.nextToken())+cumSumA[i-1];
		}
        
        // B누적합
		int m = Integer.parseInt(br.readLine());
		st = new StringTokenizer(br.readLine());
		int[] cumSumB = new int[m+1];
		for (int i = 1; i <= m; i++) {
			cumSumB[i] = Integer.parseInt(st.nextToken())+cumSumB[i-1];
		}
		
        // A의 모든 부분합
		HashMap<Integer,Integer> mapA = new HashMap<Integer,Integer>();
		for (int i = 1; i <= n; i++) {
			for (int j = 0; j < i; j++) {				
				int num = cumSumA[i] - cumSumA[j];
				if(mapA.containsKey(num)) {
					mapA.put(num, mapA.get(num)+1);
				}else {
					mapA.put(num, 1);
				}
			}
		}
		
        // B의 모든 부분합
		HashMap<Integer,Integer> mapB = new HashMap<Integer,Integer>();
		for (int i = 1; i <= m; i++) {
			for (int j = 0; j < i; j++) {				
				int num = cumSumB[i] - cumSumB[j];
				if(mapB.containsKey(num)) {
					mapB.put(num, mapB.get(num)+1);
				}else {
					mapB.put(num, 1);
				}
			}
		}

		long answer = 0;
		for(int keyA : mapA.keySet()) {
			if(mapB.containsKey(T-keyA)) { // contains대신 이분탐색을 사용하면 시간복잡도를 줄일 수 있습니다.
				answer += 1.0*mapA.get(keyA) * mapB.get(T-keyA); // int범위를 초과할 수 있다는걸 조심해야 합니다.
			}
		}
//		System.out.println(mapA.toString());
//		System.out.println(mapB.toString());
		System.out.println(answer);
	}
	
}