【TypeScript】KDTreeで最近傍探索

KDTree

k次元ユークリッド空間中の点を分割するデータ構造であり、近傍探索に非常に有効です。今回は平面（k=2）の場合で実装を行います。

KDTreeを用いた近傍探索アルゴリズム

探索木の深さに応じてxまたはy座標を閾値と比較して二分探索の要領で走査します。x,y座標の大小比較だけで済むのは普通にユークリッド距離が計算できるからですが、少し一般化すれば以下のような疑似コードで表現できます。

function kd-search(root: Node){
  var result = [] // 目的の近傍点のリスト
  search(root, result)
}

function search(node: Node, result: Array<Point>){
  result.push(node.point)
  result.sort // 距離計算
  var next1 = (最初に探索する子ノード) // node.left or node.right 
  search(next1, result)
  var next2 = (先ほど探索しなかった方の子ノード)
  if (next2に存在するかもしれない点との距離の最小値) <= (resultの距離の最大値) {
    search(next2, result)
  }
}

ポイント

• より近傍点がありそうな方を(最初に探索する子ノード)として上手く選択する
  next1を上手く選択できればnext2の方の探索が不要となりより高速な探索が可能となります。
• (next2に存在するかもしれない点との距離の最小値)を適切に計算する
  厳密な最小値を計算するのが難しい場合は、最小値以下の値（下限）で代用できます。ただし、適当に小さい値を渡すと不要な探索が増え計算時間が増大するので注意。

実装

KDTree の構築

平面$(x,y)$上の点を木の深さに応じてxまたはy座標の大小で繰り返し二分割することでKDTreeを構築します。今回は話を単純にするため、すべてのデータが事前に与えられると仮定して平衡な二分木を最初につくっておきます。動的なデータに対応するには適切な平行二分探索木の実装が別途必要です。

interface Point2D {
  x: number
  y: number
} 

/**
 * 頂点自身がひとつの座標点をもつ
 * - left には自身より小さい座標点
 * - right には自身より大きい座標点が含まれる  
 * 座標点 Point2D の大小はは頂点の depth によって x,y 座標の大小で比較する
 */
interface SearchNode extends Point2D {
  depth: number
  left?: SearchNode
  right?: SearchNode
}

function getComparator(depth: number): Comparator<Point2D> {
  if (depth % 2 === 0) {
    return (a, b) => (a.x - b.x)
  } else {
    return (a, b) => (a.y - b.y)
  }
}

function buildTree(points: Point2D[]): SearchNode {
  return buildSubTree(points, 0)
}

function buildSubTree(points: Point2D[], depth: number): SearchNode {
  points.sort(getComparator(depth))
  const mid = Math.floor(points.length / 2)
  const p = points[mid]
  cons n: SearchNode = {
    x: p.x,
    y: p.y,
    depth: depth
  }
  if (mid > 0) {
    n.left = buildSubTree(points.slice(0, mid), depth + 1)
  }
  if (mid + 1 < points.length) {
    n.right = buildSubTree(points.slice(mid + 1, points.length), depth + 1)
  }
  return n
}

KDTree による探索

interface MeasuredPoint extends Point2D {
  dist: number // 探索の中心点からの距離
}

interface EuclideanSearch {
  node?: SearchNode // 現在の探索対象
  query: Point2D // 探索の中心点
  k: number // k個の近傍点を探す
  result: MeasuredPoint[] // 近い順にソート
}

function searchEuclidean(state: EuclideanSearch) {
  const node = state.node
  if (!node) return
  const pos = state.query
  const d = measure(pos, node) // 普通のユークリッド距離
  state.result.push({
    x: node.x,
    y: node.y,
    dist: d,
  })
  state.result.sort((a,b) => a.dist - b.dist)
  while (state.result.length > k) state.result.pop()
  
  // 深さに応じてxまたはy座標で大小比較
  const compareX = (node.depth % 2 === 0)
  const threshold = compareX ? node.x : node.y
  const value = compareX ? pos.x : pos.y
  searchEuclidean({
    ...state,
    node: value < threshold ? node.left : node.right
  })
  // ここでは最小値を直接計算せず下限で代用
  // (探索しなかった方に存在するかもしれない点と点posの距離の最小値) 
  //   >= (点posと閾値thresholdが表す直線との距離)
  const dist2th = Math.abs(value - threshold)
  if (dist2th <= state.result[state.result.length - 1].dist) {
    searchEuclidean({
      ...state,
      node: value < threshold ? node.right : node.left
    })
  }
}