

 This page uses content from Wikipedia. The original article was at K-d tree. The list of authors can be seen in the page history. As with Rosetta Code, the text of Wikipedia is available under the GNU FDL. (See links for details on variance)

A k-d tree (short for k-dimensional tree) is a space-partitioning data structure for organizing points in a k-dimensional space. k-d trees are a useful data structure for several applications, such as searches involving a multidimensional search key (e.g. range searches and nearest neighbor searches).
k-d trees are a special case of binary space partitioning trees.

k-d trees are not suitable, however, for efficiently finding the nearest neighbor in high dimensional spaces. As a general rule, if the dimensionality is k, the number of points in the data, N, should be N ≫ 2k.
Otherwise, when k-d trees are used with high-dimensional data, most of the points in the tree will be evaluated and the efficiency is no better than exhaustive search, and other methods such as approximate nearest-neighbor are used instead.

Task: Construct a k-d tree and perform a nearest neighbor search for two example data sets:

 The Wikipedia example data of [(2,3), (5,4), (9,6), (4,7), (8,1), (7,2)].
 1000 3-d points uniformly distributed in a 3-d cube.
For the Wikipedia example, find the nearest neighbor to point (9, 2)
For the random data, pick a random location and find the nearest neighbor.

In addition, instrument your code to count the number of nodes visited in the nearest neighbor search. Count a node as visited if any field of it is accessed.

Output should show the point searched for, the point found,
the distance to the point, and the number of nodes visited.

There are variant algorithms for constructing the tree.
You can use a simple median strategy or implement something more efficient.
Variants of the nearest neighbor search include nearest N neighbors, approximate nearest neighbor, and range searches.
You do not have to implement these.
The requirement for this task is specifically the nearest single neighbor.
Also there are algorithms for inserting, deleting, and balancing k-d trees.
These are also not required for the task.

type
 
  Point[Dim: static Natural; T: SomeNumber] = array[Dim, T]
 
  KdNode[Dim: static Natural; T: SomeNumber] = ref object
    x: Point[Dim, T]
    left, right: KdNode[Dim, T]
 
 
func toKdNodes[N, Dim: static Natural; T](a: array[N, array[Dim, T]]): array[N, KdNode[Dim, T]] =
  ## Create an array of KdNodes from an array of list of numerical values.
  for i in 0..<N:
    result[i] = KdNode[Dim, T](x: a[i])
 
 
func dist(a, b: Point): Point.T =
  ## Return the squared distance between two points.
  for i in 0..<Point.Dim:
    let t = a[i] - b[i]
    result += t * t
 
 
func findMedian(nodes: openArray[KdNode]; slice: Slice; idx: Natural): int =
  ## Return the index of the median node in a list of nodes.
  ## The list is defined by the full list of nodes and an slice.
 
  var first = slice.a
  var last = slice.b
  if last < first: return -1
  if last == first: return first
 
  let md = first + (last - first + 1) div 2
 
  while true:
    let pivot = nodes[md].x[idx]
 
    swap nodes[md].x, nodes[last].x
    var store = first
    for i in first..last:
      if nodes[i].x[idx] < pivot:
        if i != store:
          swap nodes[i].x, nodes[store].x
        inc store
    swap nodes[store].x, nodes[last].x
 
    if nodes[store].x[idx] == nodes[md].x[idx]:
      return md
 
    if store > md:
      last = store
    else:
      first = store
 
 
func makeTree(nodes: openArray[KdNode]; slice: Slice; idx: Natural = 0): KdNode =
  ## Build a tree from a list of nodes. Return the root of the tree.
 
  if slice.b < slice.a: return nil
 
  let n = nodes.findMedian(slice, idx)
  if n < 0: return nil
 
  let idx = (idx + 1) mod result.Dim
  nodes[n].left = nodes.makeTree(slice.a..<n, idx)
  nodes[n].right = nodes.makeTree((n + 1)..slice.b, idx)
  result = nodes[n]
 
 
func nearest(root,: KdNode; point: Point; idx: Natural;
             best: var KdNode; bestDist: var root.T; nVisited: var int) =
  ## Return the node of a tree which is the nearest to a given point.
 
  if root.isNil: return
 
  let d = dist(root.x, point)
  let dx = root.x[idx] - point[idx]
  inc nVisited
 
  if best.isNil or d < bestDist:
    bestDist = d
    best = root
 
  if bestDist == 0: return
 
  let idx = (idx + 1) mod root.Dim
 
  nearest(if dx > 0: root.left else: root.right, point, idx, best, bestDist, nVisited)
  if dx * dx >= bestDist: return
  nearest(if dx > 0: root.right else: root.left, point, idx, best, bestDist, nVisited)
 
 
#———————————————————————————————————————————————————————————————————————————————————————————————————
 
 
when isMainModule:
 
  import math, random, strformat
 
 
  proc displayResult(title: string; thisPt: Point;
                     found: KdNode; bestDist: thisPt.T; nVisited: int) =
    echo title & ':'
    echo &"  Searching for {thisPt}"
    echo &"  Found {found.x}, dist = {sqrt(float(bestDist)):.5f}"
    echo &"  Seen {nVisited} nodes."
    echo ""
 
 
  proc initRandom(point: var Point) =
    for item in point.mitems:
      item = rand(1.0)
 
 
  proc runSmallTest() =
 
    let
      wp = [[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]].toKdNodes()
      #thisPt = newPoint([9, 2])
      thisPt = Point([9, 2])
      root = wp.makeTree(0..wp.high)
 
    var
      found: KdNode[root.Dim, root.T]
      bestDist = root.T(0)
      nVisited = 0
 
    root.nearest(thisPt, 0, found, bestDist, nVisited)
    displayResult("WP tree", thisPt, found, bestDist, nVisited)
 
 
  proc runBigTest() =
 
    const N = 1_000_000
    const TestRuns = 100_000
 
    randomize()
 
    var bigTree: array[N, KdNode[3, float]]
    for node in bigTree.mitems:
      new(node)
      node.x.initRandom()
 
    let root = bigTree.makeTree(0..bigTree.high)
    var thisPt: Point[3, float]
    thisPt.initRandom()
 
    var
      found: KdNode[3, float]
      bestDist = 0.0
      nVisited = 0
 
    root.nearest(thisPt, 0, found, bestDist, nVisited)
    displayResult("Big Tree", thisPt, found, bestDist, nVisited)
 
    var sum = 0
    for _ in 0..<TestRuns:
      found = nil
      nVisited = 0
      thisPt.initRandom()
      root.nearest(thisPt, 0, found, bestDist, nVisited)
      sum += nVisited
 
    echo "Big tree:"
    echo &"  Visited {sum} nodes for {TestRuns} random searches ({sum / TestRuns:.2f} per lookup)."
 
runSmallTest()
runBigTest()