BOJ 2143 - 두 배열의 합 링크
(2023.04.10 기준 G3)

문제

배열 A와 B가 주어질 때, A의 부 배열의 합과 B의 부 배열의 합이 T가 되는 모든 쌍의 개수 출력

알고리즘

모든 가능한 합에 대해 T가 되는지 확인하지만, 누적 합과 이분 탐색으로 적당한 최적화를 해야 한다.

풀이

일단 A와 B에 대해 모든 부 배열의 합을 구해야 한다. 다만, 같은 합이 여러 번 나올 수 있다. 그러므로 한 합에 대해 몇 번 나왔는지 해시 맵을 사용하여 카운트 해주자.

map Count;
for (int i = 0; i < A.size(); i++) :
	Sum = A[i], Count[Sum] += 1;
	for (int j = i + 1; j < A.size(); j++) :
    	Sum += A[j];
        Count[Sum] += 1;

이런 느낌으로 누적 합을 이용하여 모든 합에 대해 카운트를 해주자. 각 배열의 크기는 최대 1,000 이라서 충분히 넉넉하다.

그 다음은, A의 부 배열의 모든 합을 차례대로 살펴보자. B의 부 배열의 어떤 합과 T과 되는지 확인을 해야 한다. 이 때, naive하게 차례대로 살펴보면 시간 초과가 난다.
그러므로 이분 탐색을 이용해 더하면 T가 되는 'B의 부 배열의 합'을 찾아보자.
당연히 이분 탐색을 위해서 B는 무조건 정렬을 해줘야 한다.

코드

• C++

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

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

    int T;
    cin >> T;

    int n, a, S;
    vector<int> A;
    cin >> n;
    for (int i = 0; i < n; i++){
        cin >> a;
        A.push_back(a);
    }
    map<int, ll> A_sum; // A의 부 배열의 모든 합의 나온 횟수
    for (int i = 0; i < n; i++){
        S = A[i];
        A_sum[S]++;
        for (int j = i + 1; j < n; j++){ // A[i] ~ A[j]의 합을 전부 카운트
            S += A[j];
            A_sum[S]++;
            A.push_back(S); // 가능한 합 모두 저장
        }
    }
    sort(A.begin(), A.end()); // 중복 제거
    A.erase(unique(A.begin(), A.end()), A.end());

    int m, b;
    vector<int> B;
    cin >> m;
    for (int i = 0; i < m; i++){
        cin >> b;
        B.push_back(b);
    }
    map<int, ll> B_sum; // B의 부 배열의 모든 합의 나온 횟수
    for (int i = 0; i < m; i++){
        S = B[i];
        B_sum[S]++;
        for (int j = i + 1; j < m; j++){ // B[i] ~ B[j]의 합을 전부 카운트
            S += B[j];
            B_sum[S]++;
            B.push_back(S); // 가능한 합 모두 저장
        }
    }
    sort(B.begin(), B.end()); // 중복 제거
    B.erase(unique(B.begin(), B.end()), B.end());

    ll result = 0;
    int st, en, mid;
    for (auto a: A){ // A의 부 배열의 모든 합을 차례대로 살펴보자.
        st = 0, en = B.size() - 1;
        while (st <= en){
            mid = (st + en) / 2;
            if (a + B[mid] == T){ // B의 부 배별의 어떤 합과 더해서 T가 된다면 경우의 수 더하기
                result += A_sum[a] * B_sum[B[mid]];
                break;
            }
            if (a + B[mid] < T) st = mid + 1;
            else en = mid - 1;
        }
    }
    cout << result;
}

• Python

import sys; input = sys.stdin.readline
from collections import defaultdict

T = int(input())

n = int(input())
A = list(map(int, input().split()))
A_sum = defaultdict(int) # A의 부 배열의 모든 합의 나온 횟수
for i in range(n):
    S = A[i]
    A_sum[S] += 1
    for j in range(i + 1, n): # A[i] ~ A[j]의 합을 전부 카운트
        S += A[j]
        A_sum[S] += 1

m = int(input())
B = list(map(int, input().split()))
B_sum = defaultdict(int) # B의 부 배열의 모든 합의 나온 횟수
for i in range(m):
    S = B[i]
    B_sum[S] += 1
    for j in range(i + 1, m): # B[i] ~ B[j]의 합을 전부 카운트
        S += B[j]
        B_sum[S] += 1

B = sorted(B_sum.keys()) # 이분 탐색을 위해 B의 부 배열의 모든 합 정렬
result = 0
for a in A_sum.keys(): # A의 부 배열의 모든 합을 차례대로 살펴보자.
    start = 0; end = len(B) - 1
    while start <= end:
        mid = (start + end) // 2
        if a + B[mid] == T: # B의 부 배별의 어떤 합과 더해서 T가 된다면 경우의 수 더하기
            result += A_sum[a] * B_sum[B[mid]]
            break
        if a + B[mid] < T:
            start = mid + 1
        else:
            end = mid - 1
print(result)