[C++] 2143번 두 배열의 합

https://www.acmicpc.net/problem/2143

---

 

25/01/11

 

 

이전에 북마크에 넣어두었다가 해결했던 문제다. 두 가지 해결 방법으로 해결할 수 있는데, 일단 나는 첫번째 해결 방법으로만 해결했고, 나머지 하나는 아이디어만 생각해두어서 적어보려고 한다.

---

 

문제 접근 방식:

 

 

부 배열의 합은 그냥 연속된 구간 합이다.

 

당연히 나이브하게 하면 시간초과가 날게 뻔하므로, 각각의 누적 합 배열을 $AA, BB$라고 하자.

 

이제 문제는 $T$가 주어질 때 $AA[i] - AA[j] + BB[p] - BB[q] = T$를 만족시키는 $i, j, p, q$를 구하는 경우의 수 문제로 바뀌게 되었다.

 

인덱스를 총 $4$개 뽑아야 되는데, 당연히 나이브하게 이걸 $4$개 뽑아서 조건을 만족시키는 것을 찾으려고 하면 $\mathcal{O}(N^4)$으로 시간 초과가 날 것이다.

 

그래서 각 $AA[i] - AA[j]$와 $BB[p] - BB[q]$를 전부 구한 배열 혹은 집합 $AAA, BBB$가 있다고 생각했다. 이건 $\mathcal{O}(N^2)$의 시간 복잡도가 소요된다.

 

그래서 하나를 고정한 뒤에, 예를 들어 $AAA$의 값을 뽑은 뒤(그 값을 $a$라고 해보자), $T-a$가 $BBB$에 있는지 확인하면 된다고 생각했다. 만약 있다면, ($AAA$에 있는 $a$의 개수)$\times$($BBB$에 있는 $T-a$의 개수) 만큼의 경우의 수가 전체 정답에서 더해지게 된다.

 

여기서 내가 고민했던 지점은 $AAA$와 $BBB$를 unordered_map으로 구현하면 메모리가 부족하지 않을까였는데(왜냐하면 메모리 제한이 64MB여서), 그럴 걱정은 없었다. 그냥 구현해도 메모리 제한은 넘지 않았다.

 

이 경우 $\mathcal{O}(1)$의 시간 복잡도로 각 경우의 수를 확인할 수 있다.

 

만약 $AAA$와 $BBB$가 배열(vector, array)이라면 정렬한 뒤에 이분 탐색을 통해 개수를 세면 된다.

 

정렬할 때 $\mathcal{O}(N^2 \log N)$의 시간 복잡도가 소요되고, 각 경우의 수를 확인할 때마다 $\mathcal{O}(\log N)$의 시간 복잡도가 소요된다.

---

아래는 내가 위의 접근 방식과 같이 작성한 C++ 코드이다. 더보기를 누르면 확인할 수 있다.

// 2143번 두 배열의 합
// 누적 합, 
/*
접근 방법:
누적 합 갈겨
편의 상 누적합 배열을 a, b라고 하자.
그러면 a[i] - a[j] + b[p] - b[q];
인덱스를 4개 뽑아야 해
나이브로 생각하면...
1000C2 * 1000C2 대충 O(N^4)이라서 말이 안됨.
어케 줄임
하나 고정시키고 나머지를 빠르게 계산 못함?
일단 하나 고정 시켜 그러면 O(N^2*무언가)
흠 근데 경우의 수만 구하면 되는거잖아.
일단 A따로 B따로 생각해.
O(N^2) + O(N^2)
그리고 그걸 해시맵에 저장하고 빠르게 빠르게 하면 될듯
*/
#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_map>
using namespace std;
#define fastio ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define endl '\n'
typedef long long ll;

void solve(int test_case_num){
    int T; cin >> T;
    int N; cin >> N;
    vector<int> A(N+1); // 누적 합 배열 선언
    int a; A[0] = 0;
    for (int i = 0; i < N; i++){
        cin >> a;
        A[i+1] = A[i] + a;
    }
    int M; cin >> M;
    vector<int> B(M+1); // 누적 합 배열 선언
    int b; B[0] = 0;
    for (int i = 0; i < M; i++){
        cin >> b;
        B[i+1] = B[i] + b;
    }

    unordered_map<int, int> AA, BB; // A와 B에서 나올 수 있는 모든 쌍들을 저장
    for (int i = 0; i < N; i++){
        for (int j = i+1; j < N+1; j++){ // j is always bigger than i.
            AA[A[j] - A[i]]++;
        }
    }
    for (int i = 0; i < M; i++){
        for (int j = i+1; j < M+1; j++){
            BB[B[j] - B[i]]++;
        }
    }

    ll ans = 0;
    for (const auto& pair : AA){
        auto it = BB.find(T - pair.first);
        if (it != BB.end()){
            ans += (ll)(pair.second)*(it->second);
        }
    }
    cout << ans;
    return;
}

int main(void){
    fastio;
    int T; T = 1; // cin >> T;
    for (int i = 1; i <= T; i++){
        solve(i);
    }
    return 0;
}