문제
두 배열 A, B의 부분 배열(연속)의 합 쌍 중 합이 T가 되는 경우의 수를 구하라.
입력
첫째 줄에 T, 둘째 줄에 N과 A 배열, 다음 줄에 M과 B 배열이 주어진다.
출력
부분 배열 합의 쌍으로 T를 만드는 경우의 수를 출력한다.
예제
| 입력 | 출력 |
|---|---|
5 4 1 3 1 2 3 1 3 2 | 7 |
풀이
A의 모든 부분 배열 합을 HashMap에 저장한 뒤, B의 부분 배열 합과 매칭한다.
- A의 모든 연속 부분 배열 합을 구해 HashMap에 빈도를 저장한다
- B의 모든 연속 부분 배열 합 tmp에 대해
T - tmp가 HashMap에 있으면 빈도를 ans에 더한다 - 최종 ans를 출력한다
핵심 아이디어: O(N^2)개의 A 부분합을 HashMap에 넣고, O(M^2)개의 B 부분합에 대해 O(1) 조회하므로 총 O(N^2 + M^2)에 해결된다.
코드
import java.io.*;
import java.util.*;
public class Day355BOJ2143두배열의합 {
static int T, N, M;
static long ans = 0L;
static int[] A, B;
static HashMap<Integer, Integer> map;
static StringTokenizer st1, st2;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
T = Integer.parseInt(br.readLine());
N = Integer.parseInt(br.readLine());
st1 = new StringTokenizer(br.readLine());
M = Integer.parseInt(br.readLine());
st2 = new StringTokenizer(br.readLine());
A = new int[N];
for (int i = 0; i < N; i++)
A[i] = Integer.parseInt(st1.nextToken());
B = new int[M];
for (int i = 0; i < M; i++)
B[i] = Integer.parseInt(st2.nextToken());
map = new HashMap<>();
for (int i = 0; i < N; i++) {
int tmp = 0;
for (int j = i; j < N; j++) {
tmp += A[j];
Integer cnt = map.get(tmp);
if (cnt == null)
map.put(tmp, 1);
else
map.replace(tmp, cnt + 1);
}
}
for (int i = 0; i < M; i++) {
int tmp = 0;
for (int j = i; j < M; j++) {
tmp += B[j];
int dif = T - tmp;
Integer cnt = map.get(dif);
if (cnt != null)
ans += cnt;
}
}
System.out.println(ans);
br.close();
}
}복잡도
- 시간: O(N^2 + M^2) - A, B 부분합 생성 및 매칭
- 공간: O(N^2) - HashMap 저장