📖 백준 2143번 : https://www.acmicpc.net/problem/2143
조건
| 시간 제한 | 메모리 제한 |
|---|---|
| 2 초 | 64 MB |
문제
한 배열 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을 출력한다.
풀이
단순하게 생각하면 A배열이 가질 수 있는 모든 부 배열 쌍과 B배열이 가질 수 있는 모든 부 배열 쌍을 전부 비교해줘서 T를 찾는 방법을 떠올릴 수 있다. 하지만 이 방식으로는 O(n^2*m^2)의 시간 복잡도를 가져서 주어진 제한시간 내에 찾을 수 없다. 각 배열의 부 배열 쌍마다 이루는 모든 조합을 일일이 찾을 순 없으니, 어떻게 하면 빠르게 매칭시킬 수 있을지 생각해보면 이분 탐색을 떠올릴 수 있다.
각 배열의 부 배열 쌍을 구간합을 이용해서 구해 새로운 두 배열에 저장한다.
1번을 포함하는 모든 부배열의 쌍 = n
2번을 포함하고 1번을 포함하지 않는 모든 부배열의 쌍 = n-1
3번을 포함하고 1~2번을 포함하지 않는 모든 부배열의 쌍 = n-2
...
n번을 포함하고 1~n-1을 포함하지 않는 모든 부배열의 쌍 = 1
이므로, n(n+1)/2가 각 배열의 최대 크기가 된다. 나이브하게 n^2이라 생각하고 백만으로 잡아뒀다.
각 배열의 부 배열 쌍을 모두 구했으면 정렬해서 이분탐색이 가능하게 만들어준다. 만든 한 배열을 순회 하면서 나머지 한 배열에서 T를 만들 수 있는 수를 이분탐색으로 찾으면 O(n^2log(n^2))으로 답을 찾을 수 있다.
A배열의 부배열 쌍을 arr, B배열의 부배열 쌍을 brr에 저장했을 때,
정렬된 brr 배열에서 “(T - arr[i]) 이상”이 시작되는 지점과 “(T - arr[i] + 1) 이상”이 시작되는 지점을 각각 lower_bound로 찾으면, 그 사이 거리만큼이 곧 “(T - arr[i])” 값의 등장 횟수가 된다. 이 방식으로 arr[i]마다 필요한 brr의 원소 개수를 구해 다 더해주면 답이 된다.
주의해야 하는 점은 답이 int의 범위를 넘어갈 수 있기 때문에 long long int와 같은 큰 수를 담을 수 있는 자료형을 사용해야한다.
코드
#include <iostream>
#include <algorithm>
using namespace std;
int psumA[1001], psumB[1001];
int arr[1000001], brr[1000001];
int main() {
ios::sync_with_stdio(false); cin.tie(nullptr);
int T, n, m;
cin >> T;
cin >> n;
for (int i = 1; i <= n; i++) {
cin >> psumA[i];
psumA[i] += psumA[i - 1];
}
cin >> m;
for (int i = 1; i <= m; i++) {
cin >> psumB[i];
psumB[i] += psumB[i - 1];
}
int sa = 0;
for (int i = 0; i < n; i++) {
for (int j = i + 1; j <= n; j++) {
arr[sa++] = psumA[j] - psumA[i];
}
}
int sb = 0;
for (int i = 0; i < m; i++) {
for (int j = i + 1; j <= m; j++) {
brr[sb++] = psumB[j] - psumB[i];
}
}
sort(arr, arr + sa);
sort(brr, brr + sb);
long long int ans = 0;
for (int i = 0; i < sa; i++) {
auto it = lower_bound(brr, brr + sb, T - arr[i]);
ans += lower_bound(brr, brr + sb, T - arr[i] + 1) - it;
}
cout << ans;
return 0;
}
