[BOJ 2143 // C++] 두 배열의 합

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

 

이번에 볼 문제는 백준 2143번 문제인 두 배열의 합이다.
문제는 아래 링크를 확인하자.

www.acmicpc.net/problem/2143

2143번: 두 배열의 합

각 배열 A, B의 크기는 최대 1000으로 그리 큰 숫자가 아니라는 것에 주목하자.

이 문제에 사용되는 연산은 대부분 간단한 덧셈의 반복이므로, 모든 A의 구간합과 B의 구간합을 구해 같은 쌍을 찾아주어도 무리가 없다.

 

1~n, 2~n, 3~n, ... (n-1)~n, n~n의 범위의 합을 각각 저장해두면, index n+1 이하 범위에서 나올 수 있는 구간합은 n+1번째 원소 자신과, 구해진 앞쪽 범위의 n을 포함하는 구간합에 n+1번째 원소를 각각 더한 것들으로 전부 구할 수 있다. 이를 이용하여 구간합의 배열을 각각 만든다. 이 때, B의 구간합을 T에서 뺀 값을 저장한다면, 이 문제는 같은 값들의 쌍을 구하는 문제로 간단히 바뀐다.

 

그 다음으로, 두 배열을 정렬한다. 마지막으로 각 배열의 index를 가리키는 두 포인터 ptA, ptB를 만들어, index 0에서부터 두 배열을 다음과 같이 탐색한다.

 

1) ptA가 가리키는 A쪽 구간합배열값이 ptB가 가리키는 B쪽 구간합배열값보다 작다면 ptA++, 크다면 ptB++

2) 둘이 같다면, A쪽 구간합배열에 그 수가 안 나올 때까지 ptA++을 하며 A쪽 구간합배열에 그 수가 들어있는 개수 cntA를 구한다. B쪽도 마찬가지로 ptB++를 하며 B쪽 구간합배열에 그 수가 들어있는 개수 cntB를 구한다. 출력할 답 변수에 cntA * cntB를 더해준다. [cntA * cntB의 값이 int 범위를 넘어갈 수 있음에 유의하자.]

3) ptA 또는 ptB가 각 배열의 범위를 넘어설 수 있다는 것을 유의하며 위 과정을 반복한다.

 

아래는 제출한 소스코드이다.

#include <iostream>
#include <algorithm>
using std::cin;
using std::cout;
using std::sort;

int A[1000];
int B[1000];
int prefixsumA[1000000];
int prefixsumB[1000000];

int main()
{
    std::ios::sync_with_stdio(0);
    cin.tie(0);

    int T;
    cin >> T;
    int N;
    cin >> N;
    int index1 = 0;
    for (int i = 0;i < N;i++) { // A의 구간합 배열 만들기
        int x;
        cin >> x;
        for (int j = 0;j <= i;j++) {
            A[j] += x;
            prefixsumA[index1] = A[j];
            index1++;
        }
    }
    int index2 = 0;
    cin >> N;
    for (int i = 0;i < N;i++) { // B의 (T - 구간합) 배열 만들기
        int x;
        cin >> x;
        for (int j = 0;j <= i;j++) {
            B[j] += x;
            prefixsumB[index2] = T - B[j];
            index2++;
        }
    }
    sort(prefixsumA, prefixsumA + index1); // 정렬
    sort(prefixsumB, prefixsumB + index2);
    int ptA = 0; // Two Pointer
    int ptB = 0;
    long long ans = 0;
    while (ptA < index1 and ptB < index2) {
        int AA = prefixsumA[ptA];
        int BB = prefixsumB[ptB];
        if (AA < BB) ptA++;
        else if (AA > BB) ptB++;
        else {
            long long cntA = 0;
            long long cntB = 0;
            while (prefixsumA[ptA] == AA) {
                ptA++;
                cntA++;
                if (ptA >= index1) break;
            }
            while (prefixsumB[ptB] == BB) {
                ptB++;
                cntB++;
                if (ptB >= index2) break;
            }
            ans += cntA * cntB;
        }
    }

    cout << ans;

    return 0;
}