[BOJ] #2143 두 배열의 합

시간 제한	메모리 제한	정답 비율
2 초	64MB	27.919 %

2143번: 두 배열의 합

 

문제

한 배열 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가지 경우가 있다.

 

해결

key point, 누적합 + 투포인터. 자료형은 long long!!

1. A 배열과 B 배열의 누적합을 구한다. → apSum, bpSum
2. 누적합을 이용하여 가능한 부분합의 조합을 구하고 정렬한다. → A, B
3. 38번줄 부터 A와 B의 부배열 합이 T가 되는 수를 투포인터를 이용해서 구한다.
   1. 가장 왼쪽 포인터 l, 가장 오른쪽 포인터 r이 있다. l은 A배열의 가장 작은 원소를 가리키고, r은 B배열의 가장 큰 값을 가리키고 있다.
   2. A[l] + B[r] = sum 이다.
   3. 현재 sum이 T보다
      1. 크다면, sum을 줄여야 하므로 r을 왼쪽으로 한 칸 당겨 B[r]의 값을 줄여야 한다.
      2. 작다면, sum을 키워야 하므로 l을 오른쪽으로 한 칸 당겨 A[l]의 값을 키워야 한다.
      3. 같다면, 
         1. A[l]와 같은 값을 가진 원소를 배열 A에서 몇 개 있는지 찾고 → cntA
         2. B[r]와 같은 값을 가진 원소를 배열 B에서 몇 개 있는 지 찾아 → cntB
         3. cntA와 cntB를 곱해준 값이 현재 T를 만드는 경우의 수이다.
         4. 또 다른 A[l]와 B[r]의 조합으로 T를 찾기 위해서 sum을 A[l] + B[r]로 갱신시켜주고 반복문을 계속한다.
4. 반복문을 나가고난 후 ans이 정답이 된다. 

 

코드

메모리	시간
17640 KB	64 ms

#include <iostream>
#include <algorithm>
using namespace std;
int n, m;
long long t, a[1005], b[1005];
long long apSum[1005], bpSum[1005];
 
long long A[1000005], B[1000005];
 
int main() {
    scanf("%lld %d", &t, &n);
    for (int i = 1; i <= n; i++) {
        scanf("%lld", &a[i]);
        apSum[i] = apSum[i - 1] + a[i];
    }
    scanf("%d", &m);
    for (int i = 1; i <= m; i++) {
        scanf("%lld", &b[i]);
        bpSum[i] = bpSum[i - 1] + b[i];
    }
 
    int ap = 0, bp = 0;
    for (int i = 1; i <= n; i++) {
        for (int j = i; j <= n; j++) {
            A[ap++] = apSum[j] - apSum[i - 1];
        }
    }
    for (int i = 1; i <= m; i++) {
        for (int j = i; j <= m; j++) {
            B[bp++] = bpSum[j] - bpSum[i - 1];
        }
    }
    sort(A, A + ap);
    sort(B, B + bp);
 
    int l = 0, r = bp - 1;
    long long sum = A[l] + B[r], ans = 0;
    while (l < ap && r >= 0) {
        if (sum == t) {
            long long cntA = 0, cntB = 0;
            long long curA = A[l], curB = B[r];
            while (l < ap && A[l] == curA) {
                cntA++; l++;
            }
            while (r >= 0 && B[r] == curB) {
                cntB++; r--;
            }
            ans += cntA * cntB;
            sum = A[l] + B[r];
        }
        else if (sum > t) {
            sum -= B[r--]; sum += B[r];
        }
        else {
            sum -= A[l++]; sum += A[l];
        }
    }
    printf("%lld", ans);
    return 0;
}