[BOJ] 2143 - 두 배열의 합

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

BOJ - 두 배열의 합

문제 유형: 이진 탐색, 부분합, 정렬

문제 난이도: Gold III

시간 제한: 2초

메모리 제한: 64MB

 

문제

한 배열 A[1], A[2], …, A[n]에 대해서, 부 배열은 A[i], A[i+1], …, A[j-1], A[j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A[i]+…+A[j]를 의미한다. 각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

 

입력

첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A[1], …, A[n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그 다음 줄에 m개의 정수로 B[1], …, B[m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

 

출력

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

 

풀이

부분합을 이용해서 모든 부분배열의 합을 O(n^2)에 구한다.

MITM과 비슷한 논리라고 보면 된다.

A의 모든 부분배열의 합들은 정렬될 필요는 없다.

대신, B의 모든 부분배열의 합들은 정렬되어야 A가 정해지면, 해당되는 B의 요소들을 빨리 찾을 수 있다.

 

A의 모든 부분배열의 합마다, T를 만들기 위한 B의 원소의 개수를 이진탐색을 이용해서 찾는다.

즉, upper_bound() - lower_bound()를 한 것이 해당 A의조합에 대한 대응되는 B의 개수이다.

이를 누적해주면 된다.

 

수의 범위에 유의하자. long long형이여야 한다. 상한은 1,000,000으로 정해졌지만, 하한에 대한 언급은 없다.

 

코드

C++

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

int main(int argc, char *argv[]) {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  int t, n, m;
  ll a[1001], b[1001];
  vector<ll> sumA;
  vector<ll> sumB;
  cin >> t;
  cin >> n >> a[0];
  for (int i = 1; i < n; ++i) {
    cin >> a[i];
    a[i] += a[i - 1];
  }
  cin >> m >> b[0];
  for (int i = 1; i < m; ++i) {
    cin >> b[i];
    b[i] += b[i - 1];
  }

  for (int i = 0; i < n; ++i) {
    sumA.emplace_back(a[i]);
    for (int j = 0; j < i; ++j) {
      sumA.emplace_back(a[i] - a[j]);
    }
  }
  for (int i = 0; i < m; ++i) {
    sumB.emplace_back(b[i]);
    for (int j = 0; j < i; ++j) {
      sumB.emplace_back(b[i] - b[j]);
    }
  }

  sort(sumB.begin(), sumB.end());

  ll ans = 0;
  for (int i = 0; i < sumA.size(); ++i) {
    int target = t - sumA[i];
    int ub = upper_bound(sumB.begin(), sumB.end(), target) - sumB.begin();
    int lb = lower_bound(sumB.begin(), sumB.end(), target) - sumB.begin();
    ans += ub - lb;
  }

  cout << ans << "\n";

  return 0;
}