백준 2143 G3 두 배열의 합

생성 2025-04-02

해시 테이블을 사용하는 풀이

먼저, 완전 탐색으로

어렵게 생각하지 말고 무식하게 완전 탐색으로 접근해 보자.

A 배열에서, 모든 부 배열의 합과 그 갯수를 계산한다.
여기에 소요되는 복잡도는, A 배열 범위 내 인덱스를 중복하여 두 개 뽑은 후 ( $_nH_2$ ) 그 인덱스 범위를 더하는 것이므로, 대략적으로 $O(n_nH_2) = O(n^3)$ 의 복잡도이다.

B 배열도 마찬가지로 계산하면 $O(m^3)$ 정도일 것이다.

$n$ 과 $m$ 은 모두 $10^3$ 보다 작으므로, $O(n^3)$ 과 $O(m^3)$ 은 대략 1초 정도 소요되므로, 아직까지 복잡도에 문제는 없다.

A 배열의 부 배열 합 $s_a$ 과 갯수 $c_a$ 를 순회하면서, $T$ 에 도달하기 위한 값 $T-s_a$ 가 B의 부 배열의 합에 존재하는지 탐색한다.

존재한다면, 그 갯수 $c_b$ 를 $c_a$ 와 곱하여, 전체 갯수를 증가시킨다.

해시 테이블을 사용하는 이유

A 배열의 부 배열 합 $s_a$ 과 갯수 $c_a$ 를 순회하면서, $T$ 에 도달하기 위한 값 $T-s_a$ 가 B의 부 배열의 합에 존재하는지 탐색한다.

여기서 중요한 부분은 저 탐색 부분이다.

값이 존재하는지 아닌지, 탐색하는 데에 $O(1)$ 이 소요되는 아주 괜찮은 자료구조가 있는데 - 바로 해시테이블이다.

부 배열의 합과 그 합을 가지는 갯수를 저장하기 위해서 해시 테이블을 사용하면, $T$ 를 만들기 위해 필요한 값이 존재하는지 $O(1)$ 안에 확인할 수 있게 된다.

코드

코드는 아래와 같다.

#include <iostream>
#include <vector>
#include <unordered_map>

using namespace std;

int T, n, m;
vector<int> A, B;
unordered_map<int, int> sumA, sumB;

int main() {
  cin >> T;

  // A 입력
  cin >> n; 
  A.resize(n);
  for (int i = 0; i < n; ++i) {
    cin >> A[i];
  }

  // B 입력
  cin >> m;
  B.resize(m);
  for (int i = 0; i < m; ++i) {
    cin >> B[i];
  }

  // A 부 배열 합과 갯수
  for (int i = 0; i < n; ++i) {
    for (int j = i; j < n; ++j) {
      int sum = 0;
      for (int c = i; c <= j; ++c) {
        sum += A[c];
      }
      ++sumA[sum];
    }
  }

  // B 부 배열 합과 갯수
  for (int i = 0; i < m; ++i) {
    for (int j = i; j < m; ++j) {
      int sum = 0;
      for (int c = i; c <= j; ++c)  {
        sum += B[c];
      }
      ++sumB[sum];
    }
  }

  // 주의: long long
  long long count = 0;

  // A 기준으로 B에 합이 존재하는지 확인
  for (const auto& [sum_a, count_a]: sumA) {
    int required = T - sum_a;

    if (sumB.find(required) == sumB.cend()) {
      continue;
    }

    count += (long long)count_a * sumB[required];
  }

  cout << count;
}

⚠️ 주의 사항

최종 갯수의 자료형

A의 부 배열의 합의 최대 갯수는 대략 $\frac{n(n-1)}{2}$ 개, B 부 배열의 합의 최대 갯수는 대략 $\frac{m-1}{2}$ 개이다.

만일, 모든 쌍으로 $T$ 를 만들 수 있으면 가능한 총 갯수는 $\frac{n(n-1)(m)(m-1)}{4}$ 로, $10^{11}$ 에서 $10^{12}$ 사이의 값이 된다.

그리고, 이는 int형의 범위를 초과하기 때문에 최종 count를 기록하는 변수는 long long으로 설정하여 오버플로우가 발생하지 않도록 만들어야 한다.

부 배열의 합 로직 최적화

앞서 부 배열의 합을 구하는 것이 $O(n^3)$ 이라고 하였는데, 여기서 최적화를 할 수 있다.

아래가 기존의 코드이다.

for (int i = 0; i < n; ++i) { // O(n)
  for (int j = i; j < n; ++j) { // O(n)
    int sum = 0;
    for (int c = i; c <= j; ++c) { // O(n)
      sum += A[c];
    }
    ++sumA[sum];
  }
}

이 방식의 문제는 j가 1씩 증가할 때마다 i 부터 j까지의 합을 처음부터 다시 계산한다는 점에 있다. [i ... j-1]의 합을 이미 계산했는데, [i ... j]의 합을 구할 때 이 결과를 재사용하지 않고 다시 처음부터 순회하는 것이므로 비효율적이다.

아래처럼 하면 $O(n^2)$ 안에도 통과가 가능하다.

for (int i = 0; i < n; ++i) {
  int sum = 0;
  for (int j = i; j < n ; ++j) {
    // [i ... j-1]에 [j] 만 더해서, [i ... j]를 구함
    sum += A[j];
    ++sumA[sum];
  }
}

시작 인덱스 i를 고정한 상태에서, j를 증가시키면서 합을 누적해 나간다.

두 개의 루프만으로 누적 합을 구할 수 있다.

하지만, 최적화를 하지 않아도 통과 자체는 된다. 🤷‍♂️

이진 탐색을 사용하는 풀이

이진 탐색을 이용하여 풀 수도 있다.

먼저, 부 배열의 합을 무식하게 다 구하는 것 까지는 똑같다.
그런데, 각 합을 해시 테이블에 저장하는 것이 아니라, (중복된 값이 있더라도) 합의 결과를 그냥 배열에 계속 추가하자.

A의 모든 부 배열의 합을 저장한 배열을 sum_a라고 하고, B의 모든 부 배열의 합을 저장한 배열을 sum_b라고 하자.

이제, sum_a를 기준으로 삼고, sum_b를 정렬한다. sum_a의 각 원소를 기준으로, T를 만들기 위해 필요한 값 required가 sum_b에 존재하는지 이분 탐색으로 하한과 상한을 찾는다.

하한과 상한까지의 거리가 T를 만들 수 있는 값들의 개수이므로, 이를 전체 갯수에 추가한다.

시간 복잡도 비교 분석 (이진 탐색 vs 해시 테이블)

sum_b를 정렬하는 데에 $O(m^2\log{m^2})$ 이 소요되고, $n^2$ 크기의 sum_a의 각 요소에 대해, $m^2$ 크기의 정렬된 배열에 대한 이진 탐색을 수행하므로 $O(n^2\log{m^2}) = O(n^2\log{m})$ 으로, 전체 시간 복잡도는 약 $O((n^2 + m^2)\log{m})$ 이다.

해시 테이블을 사용한 방식은 required의 탐색에 $O(1)$ 이 소요되므로 전체 복잡도가 $O(n^2+m^2)$ 로 이론적으로는 더 우수하다.

하지만 이진 탐색 방식은 배열을 사용하므로 캐시 효율성 측면에서 훨씬 많은 이득을 보는 것으로 추정되고, 실제로 제출시 이진 탐색 방식이 훨씬 빠르게 통과되는 것을 알 수 있다.

코드

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int T, n, m;

vector<int> A, B;
vector<int> sum_a, sum_b;

int main() {
  cin >> T;

  // A 입력
  cin >> n;
  A.resize(n);
  for(int& e: A) {
    cin >> e;
  }

  // B 입력
  cin >> m;
  B.resize(m);
  for (int& e: B) {
    cin >> e;
  }

  // A 부 배열
  for (int i = 0; i < n; ++i) {
    int sum = 0;
    for (int j = i; j < n; ++j) {
      sum += A[j];
      sum_a.push_back(sum);
    }
  }

  // B 부 배열
  for (int i = 0; i < m; ++i) {
    int sum = 0;
    for (int j = i; j < m; ++j) {
      sum += B[j];
      sum_b.push_back(sum);
    }
  }

  // 정렬 
  sort(sum_b.begin(), sum_b.end());

  long long count = 0;

  // sum_a의 각 요소에 대해 반복:
  for (int& a: sum_a) {
    const int required = T - a;

    // 상한과 하한 찾기
    const auto& [begin, end] 
      = equal_range(
        sum_b.cbegin(), sum_b.cend(), 
        required
      );
    
    // 전체 개수 업데이트
    count += distance(begin, end);
  }

  cout << count;
}