2022년도 2회차 원티드 쇼미더코드 대회 C번 문제는 부분합 문제입니다.

두 배열 \(A, B\)에는 정수가 담겨있습니다. \(A\)의 \(k\)번째 원소 \(A_k\)와 \(B\)의 \(k\)번째 원소 \(B_k\)에 대하여,

$$ A_i + A_{i + 1} + ... + A_j = B_i + B_{i + 1} + ... + B_j \quad (i \le j) $$

를 만족하는 순서쌍 \((i, j)\)의 개수를 구하면 됩니다.

언뜻 보기에는 무척 단순해보이는 문제입니다. 풀이도 어렵지 않게 떠올릴 수 있죠. 그냥 무식하게 모든 \((i, j)\)에 대해 부분합을 구해서 비교해보면 될 것 같습니다. 이 때의 시간복잡도는 모든 \(i\)에 대해 \(n\), 그 다음 \(i\)에 대해 \(n-i\), 그리고 부분합을 구할 때 \(j-i\)이고, \(i\)와 \(j\)의 간격은 최악일 때 \(n\)이 되니까 O(\(n^3\))이 되겠네요.

물론 누적합을 활용하면 시간복잡도는 한 차원 낮출 수 있습니다. 배열에 각 원소를 담는 대신, 처음부터 누적한 합을 대신 담은 배열 \(S\)의 원소 \(S_k\)에 대해, \(i\)부터 \(j\)까지의 누적합을 계속해서 구하는 대신 \(S_j\) - \(S_i\)를 구하면 \(i+1\)부터 \(j\)까지의 부분합이 나옵니다. 이 때의 시간복잡도는 O(\(n^2\))입니다.

이 방식대로 코드를 짜면 이렇게 되겠네요.

#include <iostream>

int main()
{
    int N, answer = 0;
    std::cin >> N;
    int* A = new int[N];
    int* B = new int[N];
    for (int a = 0; a < N; ++a) std::cin >> A[a];
    for (int b = 0; b < N; ++b) std::cin >> B[b];
    for (int i = 1; i < N; ++i) {
        A[i] += A[i - 1];
        B[i] += B[i - 1];
    }
    for (int i = 0; i < N; ++i) {
        /*
         * 누적합의 차 j - i는 i + 1부터 j까지의 합만 포함합니다.
         * 0부터 i까지의 합을 비교하기 위해서는 누적합 자체에 대한 비교가 필요합니다.
         */
        if (A[i] == B[i]) ++answer;

        // 시간복잡도가 O(n ^ 2)이 되는 부분.
        // 일반적인 경우 nested for는 시간복잡도를 무조건 한 차원 올립니다.
        for (int j = i + 1; j < N; ++j) {
            if (A[j] - A[i] == B[j] - B[i]) ++answer;
        }
    }
    std::cout << answer;

    return 0;
}

그리고 이렇게 쉽게 짠 코드는 당연히 통과할 리가 없죠. 시간 초과입니다.

문제의 조건을 다시 보면, 정답은 일반적인 int의 표현범위를 벗어난다고 나옵니다. 자료형으로 long long을 쓰는 것을 넘어, 시간복잡도를 어떻게든 한 차원 더 줄여야만 하는 것이죠… 차원을 줄이기 위해서는 문제를 조금 다른 방향에서 볼 필요가 있습니다.

우선, 순서쌍을 모두 일일히 출력할 필요는 없습니다. 그렇죠? 문제를 풀 때 관심을 가지는 것은 오로지 개수입니다. 그렇다면 어느정도 불필요한 정보는 소실되어도 된다는 거죠. 음, 문제의 수식을 조금 변형시켜봅시다. 배열 \(A\)의 누적합 \(A’\)와 배열 \(B\)의 누적합 \(B’\)에 대해

$$ A'_j - A'_i = B'_j - B'_i $$

를 만족할 때 \(i+1\)부터 \(j\)까지의 부분합의 크기가 같습니다. 값을 넘겨주면,

$$ A'_j - B'_j = A'_i - B'_i $$

이 되네요.

한 배열 안에서의 누적합끼리의 차에서 두 배열의 누적합 간의 차로 바뀌었습니다. 그리고 배열 간 누적합의 차를 이용한다면, 부분합의 크기가 같은 순서쌍을 찾는 절차는 조금 달라집니다.

배열 \(D\)의 원소 \(D_k\)가 \(A’_k - B’_k\)라 할 때,

$$ D_i = D_j $$

이면 \(i+1\)부터 \(j\)까지의 부분합은 그 크기가 같습니다. 크기가 같은 원소들의 순서쌍을 찾는 문제로 바뀌었습니다. 하지만 아직까지는 이 방법으로는 시간복잡도는 여전히 O(\(n^2\))입니다.

하지만 문제에서 요구하는 답은 순서쌍의 개수입니다. 이렇게 생각해보면 어떨까요? 누적합 간의 차 \(D\)의 인덱스와는 상관없이, 그 크기에 따라 분류해서 통에 따로 담는 겁니다. 그리고 통에서 임의로 두 개를 꺼내면 되는 거죠. 크기가 같은 원소끼리 모아 그 수를 담은 배열 \(d\)를 찾습니다. 예를 들어 \(D= \lbrace 1, 1, 1, 2, 2, 3 \rbrace \)은 \(d= \lbrace 3,2,1 \rbrace \)이 됩니다.

\(n(i,j)\)를 순서쌍의 개수라 하면

$$ n(i, j)=\sum_{}^{}{_{d_k}\mathrm{C}_{2}} $$

이 되겠네요.

정확한 \(i\)와 \(j\)의 값은 소실되지만, 개수는 여전히 유효합니다. 그리고 시간복잡도는 O(\(n\))이 됩니다.

#include <iostream>
#include <cstdio>
#include <map>

int main() {
    long long N, answer = 0;
    std::cin >> N;
    std::map<long long, long long> _map; // 누적합 간의 차를 담을 배열 D. 편의상 map으로 구현
    long long* A = new long long[N];
    long long* B = new long long[N];
    
    for (long long a = 0; a < N; ++a) scanf("%lld", A + a);
    for (long long b = 0; b < N; ++b) scanf("%lld", B + b);
    _map[A[0] - B[0]] = 1;
    for (long long i = 1; i < N; ++i) { // 누적합 간의 차를 구합니다.
        A[i] += A[i - 1];
        B[i] += B[i - 1];
        if (_map.find(A[i] - B[i]) == _map.end()) _map[A[i] - B[i]] = 1;
        else _map[A[i] - B[i]] += 1;
    }
    for (const auto& [k, v] : _map) {
        /*
         * 누적합 간의 차는 어디까지나 i + 1부터 j까지의 합만을 나타냅니다.
         * 0부터 j까지의 누적합은 따로 계산을 해줘야 합니다.
         * 누적합의 차가 0일 때에는, 0을 두 개 고를 수도 있지만, 그 자체로도 부분합이 같습니다.
         * 달리 생각해보자면, 아무 것도 누적되지 않은 상태를 맨 앞에 둘 수 있습니다.
         * 그렇게 되면 0의 개수가 하나 더 많아지므로, 이 때에는
         * nC2가 아니라 (n + 1)C2가 됩니다.
         */
        if (k == 0)
            answer += v * (v + 1) / 2;
        else answer += v * (v - 1) / 2;
    }
    
    std::cout << answer;

    return 0;
}

배열의 인덱스 \(i, j\)를 각자 조회하고 그 사이의 부분합을 구하는 과정이 O(\(n^3\)).

부분합을 그때그때 구하는 대신 누적합을 미리 구해두면 그 과정은 O(\(n^2\)).

그리고, 누적합 간 차를 이용하면 그 과정은 O(\(n\))까지 줄어듭니다.

이걸 일종의 동적 프로그래밍으로 볼 수 있을까요? 잘 모르겠군요.