[백준 / BOJ] C++ 2143 두 배열의 합

2143번: 두 배열의 합

 

문제

 

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

2143번: 두 배열의 합

 

 

풀이

n과 m이 최대 1000이므로 가능한 모든 부분합을 구한 다음 가능한 조합의 수를 구하면 된다. 즉, A 부분합 배열의 각 원소에 대해 T가 되기 위해 필요한 수가 B 부분합 배열에 몇 개가 있는지 확인한다. B 부분합 배열을 오름차순 정렬한 후 lower_bound, upper_bound를 사용해 iterator의 차이(개수)를 정답에 추가한다. 경우의 수는 int 범위를 벗어날 수 있으므로 long long으로 해야 한다.

 

 

코드

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

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  int T, n, m;
  int A[1001], B[1001];
  cin >> T >> n;
  for (int i = 0; i < n; i++)
    cin >> A[i];
  cin >> m;
  for (int i = 0; i < m; i++)
    cin >> B[i];
  vector<int> aSum, bSum;
  for (int i = 0; i < n; i++) {
    int sum = A[i];
    aSum.push_back(sum);
    for (int j = i + 1; j < n; j++) {
      sum += A[j];
      aSum.push_back(sum);
    }
  }
  for (int i = 0; i < m; i++) {
    int sum = B[i];
    bSum.push_back(sum);
    for (int j = i + 1; j < m; j++) {
      sum += B[j];
      bSum.push_back(sum);
    }
  }
  sort(bSum.begin(), bSum.end());
  long long ans = 0;
  for (int i = 0; i < aSum.size(); i++) {
    int tmp = T - aSum[i];
    int u = upper_bound(bSum.begin(), bSum.end(), tmp) - bSum.begin();
    int l = lower_bound(bSum.begin(), bSum.end(), tmp) - bSum.begin();
    ans += u - l;
  }
  cout << ans;
  return 0;
}