백준 20176 Needle

 

출처 : https://www.acmicpc.net/problem/20176

 

참고한 글 : https://koosaga.com/263

 

문제 및 풀이

 

3개의 선분이 위에서부터 차례대로 하나씩 있고, 각 선분 당 특정 좌표에는 구멍이 뚫려있습니다.

아래 그림은 정답이 될 수 있는 바늘이 3개의 선분을 통과한 사례입니다. 

 

기본적으로 바늘이 3개의 선분을 지나가기 위해서는 a~b, b~c 를 지나갈 때 선분의 기울기가 같아야 하고, b의 위치가 동일해야 합니다. 즉, a+c = 2 * b 를 만족한다면 바늘은 선분을 지날 수 있습니다. (편의상, 배열 A, B, C 의 어떤 한 원소를 a, b, c 라 하겠습니다)

 

일반적으로 모든 a+c 를 구하기 위해선, 어쩔 수 없이 배열 A, C의 원소를 모두 순회하며 구해주는 수밖에 없는데 

이 문제는 각 배열의 크기가 최대 6만이므로 그런 방식으로는 시간 초과가 발생하게 됩니다.

 

따라서 a+c 를 빠르게 구하기 위해 convolution 을 활용하게 되는데, 이 부분은 fft를 사용합니다.

fft에 대한 설명은 여기 블로그에 잘 되어 있습니다.

 

구하는 방법은 만약 A원소가 1, 3, 5 이고, C원소가 2, 4 라고 하면 원소들 중 가장 큰 원소를 배열의 크기로 잡고

v[i] = i가 원소에 있으면 1, 없으면 0 으로 배열을 다시 만들어줍니다.

따라서 a는 [0, 1, 0, 1, 0, 1], b는 [0, 0, 1, 0, 1, 0] 이 됩니다. 이 둘을 벡터로 보고 convolution을 하게 되면

[0, 0, 0, 1, 0, 2, 0, 2, 0, 1] (x^9 + 2*x^7 + 2*x^5 + x^3) 이 되고 이 벡터는 a+c를 하면 등장하는 수의 횟수를 의미합니다.

따라서 9는 1회(5+4), 7은 2회(3+4, 5+2), 5는 2회(1+4, 3+2), 3은 1회 (1+2) 가 되는 것이죠.

 

이 벡터를 fft를 이용하면 O(nlgn) 으로 구할 수 있습니다.

 

따라서 정리하자면 a,c 의 convolution을 미리 구해둔 배열 v를 이용해 B원소와 비교하면 됩니다.

 

식은 다음처럼 나오는데, 한 번 더 해석하자면 선분 B에서 뚫려있는 좌표 i과 a+c를 해서 나오는 2*i가 있다면 바늘이 통과하는 경우이므로

둘을 곱해주면 됩니다. 이 과정을 B는 0~6만 까지 고정시키면서 돌리고 모두 더해주면 됩니다.

 

 

소스 코드(.cpp)

fft 코드는 koosaga님 팀노트를 참고하였습니다.

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

typedef long long ll;
typedef complex<double> base;

void fft(vector<base> &a, bool inv){
  int n = a.size(), j = 0;
  vector<base> roots(n/2);
  for(int i=1; i<n; i++){
    int bit = (n >> 1);
    while (j >= bit) {
        j -= bit;
        bit >>= 1; 
    }
    j += bit;
    if(i < j) swap(a[i], a[j]);
  }
  double ang = 2 * acos(-1) / n * (inv ? -1 : 1);
  for(int i=0; i<n/2; i++){
    roots[i] = base(cos(ang * i), sin(ang * i));
  }
  
  for(int i=2; i<=n; i<<=1){
  	int step = n / i;
    for(int j=0; j<n; j+=i){
      for(int k=0; k<i/2; k++){
        base u = a[j+k], v = a[j+k+i/2] * roots[step * k];
        a[j+k] = u+v;
        a[j+k+i/2] = u-v;
} }
}
  if(inv) for(int i=0; i<n; i++) a[i] /= n; // skip for OR convolution.
}
vector<int> multiply(vector<int> &v, vector<int> &w){
  vector<base> fv(v.begin(), v.end()), fw(w.begin(), w.end());
  int n = 2; while(n < v.size() + w.size()) n <<= 1;
  fv.resize(n); fw.resize(n);
  fft(fv, 0); fft(fw, 0);
  for(int i=0; i<n; i++) fv[i] *= fw[i];
  fft(fv, 1);
  vector<int> ret(n);
  for(int i=0; i<n; i++) ret[i] = (int)round(fv[i].real());
  return ret;
}

int main() {
	vector<int>a(60001,0), b(60001, 0), c(60001, 0);
	int n;
	cin>>n;
	for(int i=0; i<n; i++){
		int tmp; cin>>tmp;
		tmp += 30000;
		a[tmp] = 1;
	}
	cin>>n;
	for(int i=0; i<n; i++){
		int tmp; cin>>tmp;
		tmp += 30000;
		b[tmp] = 1;
	}
	cin>>n;
	for(int i=0; i<n; i++){
		int tmp; cin>>tmp;
		tmp += 30000;
		c[tmp] = 1;
	}
	
	vector<int>ac = multiply(a, c);
	
	ll ans = 0;
	for(int i=0; i<=60000; i++){
		ans += b[i] * ac[2 * i];
	}
	cout<<ans<<"\n";
	return 0;
}