백준 11714번 Midpoint

문제 요약

2차원 평면 상에서 점 $ L+M+N $개가 주어지는데, 처음 점 $L$개 $A_1,...,A_L$는 직선 $\ell_A$에 놓여있고, 그 다음 $M$개 $B_1,...,B_M$개는 직선 $\ell_B$에 놓여있고, 그 다음 N개 $C_1,...,C_N$개는 직선 $\ell_C$에 놓여있다.

이 때, $A_i$와 $B_j$의 중점이 $C_k$가 되는 $(i,j,k)$쌍의 개수를 구하는 것이 문제에서 요구하는 것이다.

 

중요한 사항은 같은 점들이 여러개 입력될 수 있다는 점과 각 점의 좌표는 정수이며 절댓값이 $10^5$이하 라는 것이다.

풀이

일단 $\ell_A, \ell_B$의 식을 구하는 것이 가능하니 수식으로 접근을 해보자. 직선의 방정식은 아래와 같이 놓을 수 있다.

 

$$
\begin{aligned}
\ell_A:y = m_ax+b_a \\
\ell_B:y = m_bx+b_b
\end{aligned}
$$

 

위 식에서 각 기울기와 절편은 주어진 점들 중에서 두 개를 뽑으면 구할 수 있다. $C$에 대한 것은 제껴두고 어떤 점 $P(x,y)$가 주어졌을 때, 이 점을 중점으로 가지는 두 점을 $P_A(x_1,m_ax_1+b_a),\;P_B(x_2,m_bx_2+b_b)$를 각 직선 $\ell_A, \; \ell_B$에서 찾아보자.

 

$$
\begin{aligned}
\frac{1}{2}(P_A+P_B) = P \\
(x_1,m_ax_1+b_a) + (x_2,m_bx_2+b_b) &= (2x,2y)
\end{aligned}
$$

 

위 식에서 절편들을 우변으로 넘겨준 다음에 행렬로 표현하면 아래와 같다.

 

$$
\begin{aligned}
\begin{bmatrix}
x_1 + x_2 \\
m_ax_1 + m_bx_2
\end{bmatrix}
=
\begin{bmatrix}
2x \\
2y-b_a-b_b
\end{bmatrix}
\ \
\begin{bmatrix}
1 & 1 \\
m_a & m_b
\end{bmatrix}
\begin{bmatrix}
x_1 \\
x_2
\end{bmatrix}
=
\begin{bmatrix}
2x \\
2y-b_a-b_b
\end{bmatrix}
\
\end{aligned}
$$

 

2X2 행렬의 역행렬은 쉽게 구해지므로 $x_1, x_2$를 구할 수 있다. 다만, $m_a = m_b$일 때는 역행렬을 구할 수 없으므로 위 식은 두 직선 $\ell_A, \ell_B$가 평행하지 않을 때만 적용 가능하다.

일단 이게 가능할 경우의 시간복잡도를 구하자.

 

모든 $C_k$에 대해서 중점이 될 수 있는 두 점 $P_A, P_B$를 구하는 것은 상수 시간이 걸림을 위 식을 통해서 알 수 있다.

그럼 이 $P_A,P_B$가 $A,B$에 있는지 확인하는 것은 map을 이용하거나 정렬을 해둔 상태라면 각각 $O(logL), O(logM)$이 걸린다.

따라서 총 시간 복잡도는 $O(LlogL + MlogM + L(logL+logM))$이다.

 

이제 두 직선이 평행한 경우인 $m_a=m_b$일 때는 어떻게 할 지 살펴보자.

$m_a \ne m_b$인 경우에는 모든 점이 중점이 될 수 있었다.

그러나, $m_a = m_b$일 때는 중점은 $\ell_a, \ell_b$사이의 하나의 직선 위에만 존재할 수 있다. 이는 위 식에서 $m_a=m_b=m$으로 놓고 보면 쉽게 알 수 있다.

그 직선을 $\ell_m$이라고 하자.

 

$$
\ell_m:y=mx+\frac{(b_a+b_b)}{2}
$$

 

중점 $P$가 이 직선 위에만 있을 수 있다는 사실을 이용해보자. 어떤 점이 한 직선 위에 놓여있다는 사실과, 그 직선의 방정식을 알고 있으면, $P$의 $x$값 혹은 $y$값만을 알고 있어도 그 값을 식에 넣어서 그 점을 구할 수 있다.

 

이걸 어떻게 이용할까? $A_i,B_j$의 중점 $P$는 $\ell_m$에 놓여있을 수 밖에 없다.

이 때, 우리는 $A_i, B_j$의 $x$값 혹은 $y$값만을 보고도 $P$가 중점인지를 알 수 있다. 이는 바꿔 말해서 점 $C_k$가 $\ell_m$위에 있을 때, 이 점이 중점이 될 수 있는 $(A_i,B_j)$의 판단을 각 $A_i, B_j$의 $x$값만을 가지고도 알 수 있다는 것이다.

같은 말을 반복해서 한 것 같으니 이젠 이걸 어떻게 답을 구하는 데에 쓸지 보자.

 

먼저 배열 $pos_A,pos_B,pos_P$를 아래와 같이 정의하자.

 

$$
\begin{aligned}
pos_A[t] &= A에서 ; x값이 ; t인 ; 점의 ; 개수 \\
pos_B[t] &= B에서 ; x값이 ; t인 ; 점의 ; 개수 \\
pos_P[t] &= 중점 ; P의 ; x값이 ; t인 ; (A_i,B_j)의 ;개수
\end{aligned}
$$

 

그러면 $pos_P$의 값을 아래와 같은 식으로 구할 수 있다.

 

$$
pos_P[t] = \sum_i{pos_A[i]*pos_B[2t-i]}
$$

 

여기서 $i,t$는 음수일수도 있다. 위 식은 Convolution으로 FFT를 이용하면 빠르게 구할 수 있다.

 

$pos_P$를 어떻게 쓸까? 만약에 $C_k$가 $\ell_m$위에 놓여 있을 때 $(A_i, B_j)$의 중점이 $C_k$가 되는 $(i,j,k)$의 개수가 $pos_P[t]$가 된다. 즉, $C$의 모든 점들에 대해서 $\ell_m$위에 있는지 판단하고, 위에 있다면 저 값을 더해주면 된다는 것이다. 이제 문제를 풀 수 있다.

 

시간복잡도를 구하자. $pos_A, pos_B$를 구하는데에는 $O(L+M)$만큼 걸릴 것이다.

이 때 중요한 것이 $pos$배열의 크기를 얼마로 잡을 것이냐가 중요한데, 입력되는 좌표들은 전부 정수이며 절댓값이 $10^5$이하다. 따라서 중점의 $x$좌표가 될 수 있는 범위는 $[-10^5,10^5]$의 정수들이기 때문에 크기는 20만보다 좀 크게 잡아주면 된다.

이제 $pos$배열의 크기를 $R$이라고 하자. 그러면 길이가 $R$인 두 수열의 Convolution을 FFT로 구현하면 $O(RlogR)$만큼이고 이 $R$값이 20만으로 제일 크기 때문에 총 시간복잡도는 $O(RlogR)$이 된다.

 

여담이지만, 이 문제에서 FFT가 쓰이는 것을 이미 알고 들어간 상태였기 때문에 풀 수 있었던 거 같다.

 

구현 상에서 주의할 점이 몇가지 있는데 A나 B의 크기가 1일 경우 직선을 구할 수가 없기 때문에 따로 처리해줘야 한다. 크기가 1이기 때문에 그냥 모든 $(A_i,B_j)$조합을 봐주면 된다.

 

그리고, 실제 입력 좌표는 음수도 들어오기 때문에 pos배열을 만들 때 인덱스에 10만을 더해주고, FFT를 한 결과를 이용할 때도 인덱스에 주의해서 접근해야한다.

 

주어진 직선이 y축과 평행할 때는 x좌표를 이용하는 게 아니라 y좌표를 이용해야하고 역행렬을 구해서 좌표를 구할 때, 실수 오차에 신경을 써줘야 한다. 실수 오차로 진짜 개고생했다.

 

아마 좌표들 정렬하고 이분탐색하는 거보다 map을 쓰면 값을 구하는거나 실수오차에 관한 부분에 있어 좀더 편할거라고 생각한다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using cdbl = complex<double>;
using ld = long double;
using Point = pair<ll,ll>;
const ld MAX_SLOPE = 1e10;
const ll POS_SZ = 2*1e5+5;
const ld EPS = 1e-4;
int L, M, N;
vector<Point> A, B, C;
ld ma, mb, mc;
ld ba, bb, bc;

void fft(vector<cdbl> &a, bool inv) {
    int n = a.size();
    // bit reversal
    for(int i=1,j=0;i<n;++i) {
        int bit = n>>1;
        while(!((j^=bit) & bit)) bit >>=1;
        if(i<j) swap(a[i],a[j]);
    }
    for(int i=1;i<n;i<<=1) {
        double x = inv? M_PI / i : -M_PI / i;
        cdbl w = cdbl(cos(x),sin(x));
        for(int j=0;j<n;j+=i<<1) {
            cdbl p = cdbl(1,0);
            for(int k=0;k<i;++k) {
                cdbl tmp = a[i+j+k] * p;
                a[i+j+k] = a[j+k] - tmp;
                a[j+k] += tmp;
                p *= w;
            }
        }
    }
    if(inv) {
        for(int i=0;i<n;++i) a[i] /= n;
    }
}

vector<ll> multiply(vector<ll> &f, vector<ll> &g) {
    vector<cdbl> pf(f.begin(), f.end()), pg(g.begin(), g.end());
    int n = 1; while (n < f.size() + g.size()) n <<= 1;
    pf.resize(n); pg.resize(n);
    fft(pf, false); fft(pg, false);
    for (int i = 0; i < n; ++i) pf[i] *= pg[i];
    fft(pf, true);
    vector<ll> ret(n);
    for (int i = 0; i < n; ++i) {
        ret[i] = (ll)round(pf[i].real());
    }
    return ret;
}

ld getSlope(vector<Point> &line) {
  ld dx = line.back().first - line.front().first;
  ld dy = line.back().second - line.front().second;
  if(dy == 0) return 0; // parallel to x-axis
  if(dx == 0) return MAX_SLOPE; // parallel to y-axis
  return dy/dx;
}

namespace std {
  bool operator< (pair<ld,ld> a, pair<ld,ld> b) {
    if(abs(a.first - b.first) < EPS && abs(a.second - b.second) < EPS) return false;
    if(abs(a.first - b.first) < EPS) return a.second < b.second;
    return a.first < b.first;
  }
}

pair<Point, Point> getPoints(Point &p) {
  ld cx = 2.0*p.first, cy = 2.0*p.second - ba - bb;
  ld alpha = (mb * cx - cy) / (mb - ma);
  ld beta = (-ma * cx + cy) / (mb - ma);
  ld alphay = alpha*ma + ba;
  ld betay = beta*mb + bb;
  if(abs(alpha - round(alpha)) > EPS) alpha = -1e7;
  if(abs(alphay - round(alphay)) > EPS) alphay = -1e7;
  if(abs(beta - round(beta)) > EPS) beta = -1e7;
  if(abs(betay - round(betay)) > EPS) betay = -1e7;
  pair<ld,ld> ap = make_pair(round(alpha), round(alphay)), bp = make_pair(round(beta), round(betay));
  return make_pair(ap, bp);
}

int main() {
  cin.tie(nullptr); ios::sync_with_stdio(false);
  cin >> L >> M >> N; A.resize(L); B.resize(M); C.resize(N);
  for(int i=0;i<L;++i) cin >> A[i].first >> A[i].second;
  for(int i=0;i<M;++i) cin >> B[i].first >> B[i].second;
  for(int i=0;i<N;++i) cin >> C[i].first >> C[i].second;
  sort(A.begin(), A.end()); sort(B.begin(), B.end()); sort(C.begin(), C.end());
  if(A.size() == 1 || B.size() == 1) {
    ll ans = 0;
    for(int i=0;i<L;++i) {
      for(int j=0;j<M;++j) {
        ll mx = (A[i].first + B[j].first) / 2;
        ll my = (A[i].second + B[j].second) / 2;
        if((mx*2 != A[i].first + B[j].first) || (my*2 != A[i].second + B[j].second)) continue;
        if(binary_search(C.begin(), C.end(), make_pair(mx,my))) ++ans;
      }
    }
    cout << ans << '\n';
    return 0;
  }
  ma = getSlope(A); mb = getSlope(B); mc = getSlope(C);
  ba = -ma*A.front().first + A.front().second;
  bb = -mb*B.front().first + B.front().second;
  bc = -mc*C.front().first + C.front().second;
  ll ans = 0;
  if(ma == mb) {
    ld m = ma, b = (ba + bb) / 2.0;
    vector<ll> pos1(POS_SZ, 0), pos2(POS_SZ, 0);
    ll ax = A[0].first, bx = B[0].first;
    ll mx = (ax+bx) >> 1;
    if(ma == MAX_SLOPE) { // parallel to y-axis
      for(int i=0;i<L;++i) {
        pos1[A[i].second+100000]++;
      }
      for(int i=0;i<M;++i) {
        pos2[B[i].second+100000]++;
      }
    }
    else {
      for(int i=0;i<L;++i) {
        pos1[A[i].first+100000]++;
      }
      for(int i=0;i<M;++i) {
        pos2[B[i].first+100000]++;
      }
    }
    vector<ll> res = multiply(pos1, pos2);
    if(ma == MAX_SLOPE) {
      for(int i=0;i<N;++i) {
        if(C[i].first == mx) ans += res[C[i].second * 2 + 200000];
      }
    }
    else {
      for(int i=0;i<N;++i) {
        ld x = C[i].first, y = C[i].second;
        ld my = m*x + b;
        if(abs(y-my) < EPS) ans += res[C[i].first * 2 + 200000];
      }
    }
  }
  else {
    vector<pair<ld,ld>> ldA(L), ldB(M);
    for(int i=0;i<L;++i) {
      pair<ld,ld> p = make_pair(A[i].first, A[i].second);
      ldA[i] = p;
    }
    for(int i=0;i<M;++i) {
      pair<ld,ld> p = make_pair(B[i].first, B[i].second);
      ldB[i] = p;
    }
    for(int i=0;i<N;++i) {
      Point ap, bp;
      tie(ap, bp) = getPoints(C[i]);
      ll num_a = upper_bound(A.begin(), A.end(), ap) - lower_bound(A.begin(), A.end(), ap);
      ll num_b = upper_bound(B.begin(), B.end(), bp) - lower_bound(B.begin(), B.end(), bp);
      ans += num_a * num_b;
    }
  }
  cout << ans << '\n';
  return 0;
}