Day 052-Q3 | KD-Tree による最近傍探索（静的2次元点群・最近傍クエリ）

問題

2次元平面上に $N$ 個の点 $P_i = (x_i, y_i)$ が与えられる。$Q$ 個のクエリ $(qx_j, qy_j)$ に対して、クエリ点から最も近い点のユークリッド距離の平方を求めよ。

制約

パラメータ	範囲
$N$	$1 \le N \le 2 \times 10^5$
$Q$	$1 \le Q \le 10^5$
座標値	$-10^9 \le x_i, y_i \le 10^9$
時間制限	3秒

入出力例

入力例 1

出力例 1

5
2
5

(3,3) の最近傍: (1,2) または (4,5) → 距離² = 5。(7,4) の最近傍: (6,3) → 距離² = 1+1 = 2。(0,0) の最近傍: (1,2) → 距離² = 1+4 = 5。

概念図: KD-Tree の構築と枝刈り

ヒント（段階的開示）

ヒント1: 方向性

KD-Tree は点群を空間的に分割する二分木。各ノードでは深さに応じて x 軸または y 軸の中央値で分割する。最近傍探索は現在の最良距離を枝刈りに使い、平均 $O(\sqrt{N})$ で動作する。

ヒント2: アプローチ

構築: 各レベルで x/y を交互に使い、その軸でソートして中央値をノードに。深さ $O(\log N)$。
探索: クエリ点に近い側を先に探索。反対側は「軸距離² ≥ best_dist」なら枝刈り。
グローバル best: リスト [INF] でミュータブルに保持（Python の closure 制約対策）。

ヒント3: コード骨格

def build(indices, depth):
    if not indices: return -1
    if len(indices) == 1: return indices[0]
    axis = depth % 2
    indices.sort(key=lambda i: (tree_x[i] if axis==0 else tree_y[i]))
    mid = len(indices) // 2
    node = indices[mid]
    left_child[node] = build(indices[:mid], depth+1)
    right_child[node] = build(indices[mid+1:], depth+1)
    return node

def search(node, qx, qy, depth):
    if node == -1: return
    d = (tree_x[node]-qx)**2 + (tree_y[node]-qy)**2
    if d < best[0]: best[0] = d
    axis = depth % 2
    val = qx if axis==0 else qy
    pivot = tree_x[node] if axis==0 else tree_y[node]
    first = left_child[node] if val<=pivot else right_child[node]
    second = right_child[node] if val<=pivot else left_child[node]
    search(first, qx, qy, depth+1)
    if (val-pivot)**2 < best[0]:
        search(second, qx, qy, depth+1)

模範解答 (Python)

import sys
input = sys.stdin.readline

def solve():
    N, Q = map(int, input().split())
    pts = []
    for _ in range(N):
        x, y = map(int, input().split())
        pts.append((x, y))

    INF = float('inf')
    nodes = list(range(N))
    tree_x = [pts[i][0] for i in range(N)]
    tree_y = [pts[i][1] for i in range(N)]
    left_child = [-1] * N
    right_child = [-1] * N

    def build(indices, depth):
        if not indices:
            return -1
        if len(indices) == 1:
            return indices[0]
        axis = depth % 2
        indices.sort(key=lambda i: (tree_x[i] if axis == 0 else tree_y[i]))
        mid = len(indices) // 2
        node = indices[mid]
        left_child[node] = build(indices[:mid], depth + 1)
        right_child[node] = build(indices[mid+1:], depth + 1)
        return node

    root = build(nodes, 0)
    best = [INF]

    def search(node, qx, qy, depth):
        if node == -1:
            return
        nx, ny = tree_x[node], tree_y[node]
        d = (nx - qx) ** 2 + (ny - qy) ** 2
        if d < best[0]:
            best[0] = d
        axis = depth % 2
        val = qx if axis == 0 else qy
        pivot = nx if axis == 0 else ny
        if val <= pivot:
            search(left_child[node], qx, qy, depth + 1)
            if (val - pivot) ** 2 < best[0]:
                search(right_child[node], qx, qy, depth + 1)
        else:
            search(right_child[node], qx, qy, depth + 1)
            if (val - pivot) ** 2 < best[0]:
                search(left_child[node], qx, qy, depth + 1)

    sys.setrecursionlimit(300000)
    out = []
    for _ in range(Q):
        qx, qy = map(int, input().split())
        best[0] = INF
        search(root, qx, qy, 0)
        out.append(str(best[0]))

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

solve()

Step-by-Step 解説

1KD-Tree の構築
各レベルで軸（depth % 2: 0→x, 1→y）を決め、その軸の中央値でノードを分割。再帰的に左右部分木を構築。深さ $O(\log N)$、構築時間 $O(N \log N)$。

2軸による分割の意味
x 軸で分割すると、左部分木には「分割ノードより x が小さい点」が入る。これにより、クエリ点から遠い側の探索を枝刈りできる。

3枝刈りの条件
クエリ点 $(qx, qy)$ と分割軸の距離の平方 $(val - pivot)^2$ が現在の best_dist 以上なら、反対側を探索しても更新できないので枝刈り（$\ge$ でなく $<$ なら探索）。

4計算量の考察
最悪 $O(N)$ だが、ランダム点群では期待 $O(\sqrt{N})$ 程度。入力が敵対的な場合は TLE しうるため、shuffle 付き構築やランダム選択が有効。

計算量

構築: $O(N \log^2 N)$（各レベルでのソート）
最近傍クエリ: 期待 $O(\sqrt{N})$、最悪 $O(N)$
全体: $O(N \log^2 N + Q\sqrt{N})$ 期待
空間: $O(N)$

よくあるミス

ミス	原因	正しい書き方
枝刈り条件の向き	`<` と `>=` を混同	`if (val-pivot)**2 < best[0]: search(second...)`
再帰深度超過	$N=2\times10^5$ で深さ最大 $N$	`setrecursionlimit(300000)` 設定
平方根で比較	sqrt の計算コスト・精度問題	距離の平方のまま比較・出力
left/right の割り当て	`indices[mid]` をノードにし両側を正確に分割	`build(indices[:mid])` と `build(indices[mid+1:])`

次のステップ

発展問題: $k$ 近傍クエリ（最大ヒープで $k$ 個管理 + 枝刈り条件を $k$ 番目距離で制御）
関連: Range Tree（矩形範囲カウント $O(\log^2 N)$）、Ball Tree（高次元最近傍）
応用: オフライン最近傍 → 全点間最近傍ペア（$O(N \log N)$ 分割統治）

自己評価

完全理解おおむね理解要復習