2차원 평면에서 모든 점 쌍의 유클리드 거리 합 구하기

오늘은 이런 문제를 풀어보자.

 

2차원 직교 좌표계에서 N개의 점이 있다. ($N \leq 10^{6}$) 좌표가 동일한 두 점이 존재할 수 있다.

i번째 점을 $p_{i}$라 하고 i번째 점과 j번째 점의 유클리드 거리 제곱을 $dist(p_{i}, p_{j})$라고 했을 때 다음을 구하시오.

$\sum\limits^{N - 1}_{i = 1}\sum\limits^{N}_{j = i + 1} dist(p_{i}, p_{j})$

 

지금 생각할 수 있는 간단한 방법은 $\binom{N}{2}$의 쌍의 거리를 일일이 계산하는 것이다.

 

하지만 그런 걸로 되면 이 글을 쓸 필요도 없다. N이 100만까지 갈 수 있으므로 컴퓨팅으로는 빠른 계산이 당연히 어렵다.

 

이 포스트에서는 이 문제를 무려 $O(N)$에 풀어볼 것이다.

 

아이디어를 설명하기 전 다음 상황을 보자.

 

 

(0, 0)에 500개가 있고 (5, 6)에 1개가 있을 때 거리를 쉽게 계산하는 방법을 떠올려보자.

 

방법을 생각하는 건 그리 어렵지 않다. (0, 0)과 (5, 6)의 거리를 구하고 500을 곱하면 된다.

 

(5, 6)의 점을 501번째 점이라 하면 이전 500개 점의 위치는 알고 있으니 일일이 계산하지 않고 500을 곱할 수 있는 것이다.

 

아이디어는 위처럼 "지금까지 계산한 거리를 누적하면 다음 점과 이전 점까지의 거리를 O(1)에 구할 수 있지 않을까?"에서 출발한다.

 

편의를 위해 0-indexed 인덱싱을 하겠다.

 

i번째 점을 본다. 그리고 i번째 점의 x좌표를 x, i번째 점의 y좌표를 y라고 했을 때 $i \times x^{2}$와 $i \times y^{2}$를 더한다.(0-indexed이므로 i - 1이 아니라 i임을 유의한다.)

 

왜 저 값을 더하는지 이해하기 위해선 다음 완전제곱 식을 살펴봐야 한다.

 

$(a - b)^{2} = a^{2} - 2ab + b^{2}$

 

여기서 $a^{2}$가 바로 $i \times x^{2}$와 $i \times y^{2}$에 해당하는 값이다. 그리고 $b^{2}$가 바로 i - 1번째 점까지 누적된 거리를 의미한다. 그러면 각 점마다 위 식의 결과를 더해나가면 $O(1)$의 시간 복잡도로 i번째 점과 [0, i - 1]번째 점 쌍의 유클리드 거리 제곱을 구할 수 있다.

 

그러면 $2ab$와 $b^{2}$을 바로 구하기 위해 점들에 대해 순회하면서 각각 $x$, $y$, $x^{2}$와 $y^{2}$를 누적해야 한다는 것을 이해할 수 있을 것이다. 이들을 변수 sumx, sumy, sqx와 sqy에 저장하도록 하겠다.

 

변수들을 모두 정리했으니 정리하면 다음과 같이 될 것이다.

 

$ans += i \times x^{2} + i \times y^{2} + sqx + sqy - 2 \times sumx \times x - 2 \times sumy \times y$

 

여기서 $2 \times sumx \times x$와 $2 \times sumy \times y$는 sumx와 sumy에서 이미 i가 누적된 상태니 i를 또 곱할 필요가 없음에 유의한다.

 

위 계산을 끝난 다음 sumx, sumy, sqx와 sqy에 각각 값을 누적해주도록 한다.

 

이렇게 모든 점을 순회하고 나온 ans가 우리가 푸려던 문제의 답이 된다.

 

위 과정을 C++ 코드로 구현하면 다음과 같이 된다.

 

using ll = int_fast64_t;
using pll = pair<ll, ll>;

ll distance_sum_point_pairs(vector<pll> &target)
{
  ll ret = 0;
  ll sqx = 0, sqy = 0, sumx = 0, sumy = 0;

  int sz = target.size();
  for (int i = 0; i < sz; i++)
  {
    ll x = target[i].first, y = target[i].second;
    ret += sqx;
    ret -= (sumx<<1) * x;
    ret += i * x * x;

    sqx += x * x;
    sumx += x;

    ret += sqy;
    ret -= (sumy<<1) * y;
    ret += i * y * y;

    sqy += y * y;
    sumy += y;
  }

  return ret;
}

 

지금까지 2차원 평면에서 무수히 많은 점들의 유클리드 거리 쌍의 합을 구하는 방법을 알아보았다.

연습 문제

Meta Hacker Cup 2022 Round 1 B : Watering Well

 

https://www.facebook.com/codingcompetitions/hacker-cup/2022/round-1/problems/B2

 

나무들과 우물들의 좌표가 주어질 때 모든 나무<->우물 사이의 거리 제곱 합을 구하는 문제다. 우선 나무와 우물 상관없이 모든 점 쌍의 거리 제곱 합을 구한 다음 나무 쌍의 거리 제곱 합과 우물 쌍의 거리 제곱 합을 빼면 된다.