Day 128-Q2 — Kd-tree（k次元空間の最近傍探索）

2026-08-20 赤色 Master / Phase 8+ ★★★★★★★★★ 軸交互分割の静的構築+距離による枝刈りでO(log N)平均の最近傍探索

問題

2次元平面上に $N$ 個の点が与えられる。$Q$ 個のクエリ点それぞれについて、ユークリッド距離が最小となる点の番号（1-indexed）とその距離の2乗を出力せよ。

入力形式

N
x_1 y_1
...
x_N y_N
Q
qx_1 qy_1
...
qx_Q qy_Q

制約

$1 \le N, Q \le 200000$
$-10^9 \le x_i,y_i,qx_i,qy_i \le 10^9$
全点は相異なる
最近傍が複数ある場合はどれを出力してもよい

入出力例

入力例1

4
0 0
5 4
3 3
8 1
2
4 4
0 1

出力例1

2 1
1 1

(4,4)への最近傍は点2(5,4)で距離2乗1。(0,1)への最近傍は点1(0,0)で距離2乗1。

概念図

軸交互分割（depth 0: x軸 / depth 1: y軸） x=180(root分割) y分割(左) y分割(右) クエリ点との距離で 近い側を先に探索 分割線までの距離^2 >= 現在の最良距離なら 反対側は探索しない（枝刈り） 黄=分割点(中央値) / 白=葉の点

ヒント（段階的開示）

ヒント1（方向性）

クエリごとに全点との距離を計算すると$O(NQ)$で間に合わない。2次元平面を再帰的に半分ずつ区切る構造を1度だけ構築し、各クエリで不要な領域を枝刈りしながら探索するKd-treeを使う。

ヒント2（アプローチ）

深さ0はx座標、深さ1はy座標というように軸を交互に切り替えながら中央値で分割する。探索時はクエリ点が分割軸のどちら側にあるかで近そうな側を先に探索し、反対側は「分割軸までの距離^2が現在の最良距離未満の場合のみ」探索する。

ヒント3（誘導） 
def search(node, qx, qy):
    if node is None:
        return
    px, py, pid = node['point']
    d = sqdist(px, py, qx, qy)
    update_best(d, px, py, pid)
    axis = node['axis']
    diff = (qx - px) if axis == 0 else (qy - py)
    near, far = (node['left'], node['right']) if diff < 0 else (node['right'], node['left'])
    search(near, qx, qy)
    if diff * diff < best_dist:
        search(far, qx, qy)

模範解答 (Python)

import sys


def solve():
    sys.setrecursionlimit(1 << 20)
    data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    pts = []
    for i in range(n):
        x = int(data[idx]); idx += 1
        y = int(data[idx]); idx += 1
        pts.append((x, y, i + 1))

    def build(points, depth):
        if not points:
            return None
        axis = depth % 2
        points.sort(key=lambda p: p[axis])
        mid = len(points) // 2
        return {
            'point': points[mid],
            'axis': axis,
            'left': build(points[:mid], depth + 1),
            'right': build(points[mid + 1:], depth + 1),
        }

    root = build(pts, 0)

    q = int(data[idx]); idx += 1
    out = []

    def sqdist(px, py, qx, qy):
        return (px - qx) ** 2 + (py - qy) ** 2

    for _ in range(q):
        qx = int(data[idx]); idx += 1
        qy = int(data[idx]); idx += 1

        best_dist = [float('inf')]
        best_point = [None]

        def search(node):
            if node is None:
                return
            px, py, pid = node['point']
            d = sqdist(px, py, qx, qy)
            if d < best_dist[0]:
                best_dist[0] = d
                best_point[0] = (px, py, pid)

            axis = node['axis']
            diff = (qx - px) if axis == 0 else (qy - py)
            near, far = (node['left'], node['right']) if diff < 0 else (node['right'], node['left'])

            search(near)
            if diff * diff < best_dist[0]:
                search(far)

        search(root)
        px, py, pid = best_point[0]
        out.append(f"{pid} {best_dist[0]}")

    print('\n'.join(out))


solve()

Step-by-Step 解説

1軸交互分割による静的構築
深さが偶数はx座標、奇数はy座標でソートし中央値をノードに、残りを左右の子へ再帰的に渡す。木の高さは$O(\log N)$に抑えられる。
2最近傍探索の基本方針
各ノードで距離を計算し最良解を更新後、クエリ点が分割軸のどちら側かで近い側を先に探索する。
3枝刈りの正しさ
分割軸までの距離がdiffのとき、反対側のどの点も距離は少なくとも|diff|以上。diff^2 >= best_distなら反対側は探索不要。
4平均・最悪計算量
一様分布なら平均$O(\log N)$、点が偏っていると最悪$O(N)$になり得る。
5実装上の注意
クエリごとにsearchをクロージャとして定義しているため、$Q$が大きい場合は明示スタックによる反復実装で高速化できる。

よくあるミス

ミス原因正しい書き方
枝刈り条件をdiff^2 <= best_distとする等号を含めると不要な探索が増える等号なしの<で十分
分割軸の判定でaxisを使い忘れる軸交互分割の実装漏れdepth % 2で軸を切り替える
距離を浮動小数点sqrtで計算する実数演算への依存座標が整数なら距離の2乗のまま整数比較する
深い再帰でスタックオーバーフローする再帰上限未設定・未反復化setrecursionlimitを上げるか明示スタックで反復実装する

次のステップ

  • 発展: 動的な点の追加・削除が必要ならVantage-Point Treeや平方分割、定期的な再構築戦略を検討する
  • 次回予告: 離散対数問題（Baby-step Giant-step法）

自己評価

自分の回答

気づき・メモ