[백준/2143] - 두 배열의 합(JAVA)

한 배열 A[1], A[2], …, A[n]에 대해서, 부 배열은 A[i], A[i+1], …, A[j-1], A[j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A[i]+…+A[j]를 의미한다. 각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

입력

---

첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A[1], …, A[n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그 다음 줄에 m개의 정수로 B[1], …, B[m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

출력

---

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

풀이

---

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.StringTokenizer;

public class Main {
	static StringTokenizer st;
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int T = Integer.parseInt(br.readLine());
		int n = Integer.parseInt(br.readLine());
		int[] A = new int[n];
		st = new StringTokenizer(br.readLine());
		for (int i=0; i<n; i++){
			A[i] = Integer.parseInt(st.nextToken());
		}
		int m = Integer.parseInt(br.readLine());
		int[] B = new int[m];
		st = new StringTokenizer(br.readLine());
		for (int i=0; i<m; i++){
			B[i] = Integer.parseInt(st.nextToken());
		}
		Map<Long, Long> mapA = new HashMap<>();
		Map<Long, Long> mapB = new HashMap<>();
		for (int i=0; i<n; i++){
			long sum = 0;
			for (int j=i; j<n; j++){
				sum += A[j];
				long count = mapA.getOrDefault(sum, 0L);
				mapA.put(sum, count+1);
			}
		}
		for (int i=0; i<m; i++){
			long sum = 0;
			for (int j=i; j<m; j++){
				sum += B[j];
				long count = mapB.getOrDefault(sum, 0L);
				mapB.put(sum, count+1);
			}
		}
		long answer = 0;
		for (Map.Entry<Long, Long> entryA :mapA.entrySet()){
			Long aSum = entryA.getKey();
			Long aCount = entryA.getValue();
			long diff = T - aSum;
			Long bCount = mapB.getOrDefault(diff, 0L);
			answer += (aCount * bCount);
		}
		System.out.println(answer);
	}
}

해설

---

누적합 + 부배열을 다룰 수 있는지를 물어보는 문제였다.

슬라이딩 윈도우로도 풀 수 있는 문제지만, 어짜피 개수 * 개수를 리턴해야 하는 문제라고 생각해서

Map에 누적합이 몇개 있는지 count를 넣어주었다.

예를 들어 A = {1,3,1,2}, B = {1,3,2}라면

mapA에는 7가지 경우의 수(1,2,3,4,5,6,7)

mapB에는 6가지 경우의 수(1,2,3,4,5,6)가 들어갈 것이다.

 

그중에서 연속되어 있는 수열을 찾기 위해 중첩포문을 두번 돌려줄 것이다.

그리고 map의 entryset을 뒤져서

mapA = 1이라면, mapB = T-mapA = 4가 되어야 하므로

mapB의 4에 해당하는 entryset을 찾아서

mapA의 count (값이 1이 되는 부분수열의 개수) * mapB의 count(값이 2가 되는 부분수열의 개수)를 해주고 print한다.

 

물론 sliding window나 이분 탐색으로 풀 수도 있다.

그러나 규칙 / 분류에 의거해서 기계적으로 문제를 푸는 일을 굉장히 지양하려고 해서 어떻게 할까 생각하다 보니

map으로 풀어보았고, 시간도 오히려 sliding window에 비해 빨랐다.