Day 150-Q1 | kd木（2次元最近傍探索）

問題

平面上の $N$ 個の点（1-indexed、点 $i$ は座標 $(x_i, y_i)$）が与えられる。以下の構築規則に従って2次元kd木を構築し、$Q$ 個のクエリ点それぞれについて、木を辿って求めた最近傍点（ユークリッド距離の2乗が最小の点。複数該当する場合は入力インデックスが最小のもの）のインデックスを出力せよ。

構築規則（深さ $d$、$d=0$ が根）

深さ $d$ が偶数なら分割軸は $x$ 座標、奇数なら $y$ 座標。
現在のノード集合を分割軸の値で安定ソート（値が等しい場合は入力インデックス昇順）。
ソート後の下側中央（0-indexedで(サイズ-1)//2番目）をノードとする。
前方を左部分木、後方を右部分木として深さ $d+1$ で再帰構築。

入力形式

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

制約

$1 \le N \le 1000$

$1 \le Q \le 1000$

座標は $0$ 以上 $10^4$ 以下の整数

点の座標に重複がある場合がある

入出力例

入力例1

出力例1

5
2
1

概念図：構築されるkd木の構造

ヒント（段階的開示）

ヒント1（方向性）

kd木は「1次元の二分探索木」を2次元に拡張したもの。深さごとに軸を交互に切り替えながら空間を半分ずつに分割していく。最近傍探索では、まず「近そうな側」を先に見て良い候補を確保し、その後「反対側にもっと近い点があり得るか」を分割超平面までの距離だけで判定する。

ヒント2（アプローチ）

分割超平面までの距離が現在の最良距離より大きいなら、反対側のどの点も分割超平面より遠い位置にしかないため、最良距離を更新できない。分割軸の座標差の2乗と現在の最良距離2乗を比較するだけで部分木全体を探索対象から除外できるのが枝刈りの正体。

ヒント3（誘導）

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

def dist2(a, b):
    return (a[0]-b[0])**2 + (a[1]-b[1])**2

模範解答 (Python)

import sys


def build(points, depth=0):
    if not points:
        return None
    axis = depth % 2
    points = sorted(points, key=lambda p: (p[1][axis], p[0]))
    mid = (len(points) - 1) // 2
    node = points[mid]
    left = build(points[:mid], depth + 1)
    right = build(points[mid + 1:], depth + 1)
    return {"idx": node[0], "pt": node[1], "axis": axis, "left": left, "right": right}


def dist2(a, b):
    return (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2


def nearest(tree, q):
    best = [None]

    def search(t):
        if t is None:
            return
        d = dist2(t["pt"], q)
        key = (d, t["idx"])
        if best[0] is None or key < best[0]:
            best[0] = key
        axis = t["axis"]
        diff = q[axis] - t["pt"][axis]
        near, far = (t["left"], t["right"]) if diff < 0 else (t["right"], t["left"])
        search(near)
        if diff * diff <= best[0][0]:
            search(far)

    search(tree)
    return best[0][1]


def main():
    data = sys.stdin.read().split()
    ptr = 0
    n = int(data[ptr]); ptr += 1
    points = []
    for i in range(1, n + 1):
        x = int(data[ptr]); y = int(data[ptr + 1]); ptr += 2
        points.append((i, (x, y)))
    tree = build(points)
    q = int(data[ptr]); ptr += 1
    out = []
    for _ in range(q):
        qx = int(data[ptr]); qy = int(data[ptr + 1]); ptr += 2
        out.append(str(nearest(tree, (qx, qy))))
    print("\n".join(out))


main()

Step-by-Step 解説

1構築規則を機械的に適用する
深さ0では$x$座標でソートし下側中央を根にする。6点中(6-1)//2=2番目、つまり点2 (5,4) が根。

2木を明示的に確認する
root=idx2(x軸)、左=idx1(y軸、右子のみidx4)、右=idx6(y軸、左子idx5・右子idx3)という構造になる。

3最近傍探索を1クエリ追う
クエリ(9,2)は root→右(idx6, 距離2乗4)→右(idx3, 更新なし)→分割差0≤4なので左(idx5, 距離2乗2で更新)→root左部分木は分割差16>4で枝刈り。答えはidx5。

4枝刈りなしとの比較
枝刈り条件を外すと全ノードを毎回訪問し、$N$が大きいと線形探索と同じ計算量に劣化する。

よくあるミス

ミス	原因	正しい書き方
同座標点での構築結果が入力順に依存	ソートキーに軸値だけ使う	ソートキーに必ずインデックスを含める
分割差の判定を`<`のみにする	境界ケースの見落とし	等号を含む`<=`で判定する
近い側・遠い側の判定で`diff==0`の扱いが曖昧	構築規則との対応漏れ	`diff<0`なら左、それ以外(0含む)は右と明確化
深さではなく部分木の高さで軸を切り替える	再帰引数の誤り	軸は根からの深さで決まる
ブルートフォースと照合せず提出	検証不足	小さい$N$では必ずブルートフォース最近傍と一致するはず

次のステップ

発展: $k$-NNを求めるには「現在の最良候補」をヒープで$k$個保持し、枝刈り条件をヒープ内最大距離との比較に変える。
次回予告: Cycle Canceling法（最小費用流の負閉路除去法）

自己評価

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

Day 150-Q1 — kd木（k-dimensional tree・2次元最近傍探索）

問題