[Gold III] 두 배열의 합 - 2143
성능 요약
메모리: 96108 KB, 시간: 604 ms
분류
이분 탐색(binary_search), 누적 합(prefix_sum)
문제 설명
한 배열 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을 출력한다.
풀이
N, M이 1000이므로 O(N^2)로 가능한 부 배열의 값을 모두 구했다.
처음에는 부배열의 값을 모두 구할 생각을 하지 않고 누적합만 구한 후 값을 찾으려고 하였는데, 부 배열의 값을 모두 구해서 정렬해두지 않으면 풀기가 어려웠다.
list에 값을 저장한 후 정렬하여 투포인터로 구현하거나, hashmap에 부 배열의 값이 몇 번 나왔는지를 value로 저장하고 가능한 조합의 개수를 구하는 방법이 있는데, 후자의 방법이 더 쉬울 것 같아서 그렇게 구현하였다.
다른 풀이에는 이분탐색을 활용하는 방법이 있었는데, 중복된 값이 있을 수 있으니 lower bound와 upper bound를 구해서 가능한 경우의 수를 모두 구하는 접근은 생각해내기 어려울 것 같았다... 이분 탐색을 좀 더 공부해야 할듯
그리고 가능한 조합의 개수가 int 범위를 초과할 수 있기 때문에 long으로 타입을 지정해야 한다.
package BOJ.Gold.g3;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;
public class BOJ_2143_두배열의합 {
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 + 1];
StringTokenizer st = new StringTokenizer(br.readLine());
Map<Integer, Integer> sumA = new HashMap<>();
getSum(st, N, sumA, A);
int M = Integer.parseInt(br.readLine());
Map<Integer, Integer> sumB = new HashMap<>();
int[] B = new int[M + 1];
st = new StringTokenizer(br.readLine());
getSum(st, M, sumB, B); // 부배열의 합 (sumB)가 몇 번 등장하는지 hashmap에 저장
long ans = 0; // 개수가 long 범위 초과 가능
for (Integer valA : sumA.keySet()) {
int cntA = sumA.get(valA);
if (sumB.containsKey(T - valA)) {
int cntB = sumB.get(T - valA);
ans += (long) cntA * cntB;
}
}
System.out.println(ans);
}
private static void getSum(StringTokenizer st, int m, Map<Integer, Integer> sumB, int[] b) {
// 부 배열의 합을 hashmap에 저장하기
for (int i = 1; i <= m; i++) {
b[i] = b[i - 1] + Integer.parseInt(st.nextToken());
for (int j = 1; j <= i; j++) {
int sum = b[i] - b[j - 1];
if (sumB.containsKey(sum)) {
sumB.put(sum, sumB.get(sum) + 1);
} else
sumB.put(sum, 1);
}
}
}
}