7453 - 합이 0인 네 정수

문제 링크

---

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

문제 요약

---

길이가 \(N\)인 네 개의 배열이 주어진다. 이 때 네 개의 배열에서 각각 한개의 수를 뽑아 더했을 때 이 합이 0이 되는 경우의 수를 구하는 문제다.

문제 풀이

---

일단 모든 경우를 전부 보면 \(O(N^4)\)이기 때문에 시간초과가 나기 때문에 다른 방법이 필요하다. 우선 네 개의 배열 \(A\), \(B\), \(C\), \(D\)이 있다. 이 때 \(A\)와 \(B\)의 가능한 모든 합을 저장한 배열을 \(X\), \(C\)와 \(D\)의 가능한 모든 합을 저장한 배열을 \(Y\)라고 하자. 이 때 \(X\)에서 어떠한 수 한개를 뽑고, \(Y\)에서 어떠한 수 한개를 뽑아 이 수가 \(0\)이 된다고 생각을 해보자. 그렇다면 이 경우는 네 개의 배열에서 각각 한개의 수를 뽑아 0이 되는 경우와 동일하다는 것을 어렵지 않게 알 수 있다.

이제 문제는 \(X\)와 \(Y\), 두 개의 배열에서 합이 0이되는 경우의 수를 찾는 문제로 축소됐다. 이 것은 \(binary\ search\)를 통해 쉽게 구할 수 있다.

\(X\)와 \(Y\)의 배열의 크기는 \(O(N^2)\)이고 따라서 전체 시간복잡도는 \(O(N^2logN^2)\)이다.

소스 코드

---

#include <stdio.h>
#include <vector>
#include <algorithm>

using namespace std;

int n, x;
vector<int> vi[4];

vector<int> getSum(int x, int y) {
  vector<int> ret;
  for (int i = 0 ; i < (int)vi[x].size() ; ++i) {
    for (int j = 0 ; j < (int)vi[y].size() ; ++j) {
      ret.push_back(vi[x][i] + vi[y][j]);
    }
  }
  return ret;
}

int main() {
  scanf("%d", &n);
  for (int i = 0 ; i < n ; ++i) {
    for (int j = 0 ; j < 4 ; ++j) {
      scanf("%d", &x);
      vi[j].push_back(x);
    }
  }
  vector<int> ab = getSum(0, 1), cd = getSum(2, 3);
  sort(cd.begin(), cd.end());
  
  long long ans = 0;
  for (int target : ab) {
    auto lo = lower_bound(cd.begin(), cd.end(), -target);
    auto hi = upper_bound(cd.begin(), cd.end(), -target);
    ans += hi - lo;
  }
  printf("%lld\n", ans);
}