kdtree is a pure Nim k-d tree implementation. k-d trees are data structures for performing efficient spatial query operations on point data sets.

import random, strformat
import kdtree

let numPoints = 100_000
var
  points = newSeqOfCap[array[2, float]](numPoints)
  values = newSeqOfCap[int](numPoints)
  x: float
  y: float
  r = initRand(34)

for a in 0..<numPoints:
  x = r.rand(100.0)
  y = r.rand(100.0)
  points.add([x, y])
  values.add(a)

echo fmt"Building tree of {numPoints} random points..."
var tree = newKdTree[int](points, values)

# Perform nearestNeighour searches
let numSearches = 10_000
for a in 0..<numSearches:
  x = r.rand(100.0)
  y = r.rand(100.0)
  let (pt, values, dist) = tree.nearestNeighbour([x, y])
  echo fmt"point={pt}, value={value}, dist={dist}"

# Perform nearestNeighours searches
let n = 10
for a in 0..<numSearches:
  x = r.rand(100.0)
  y = r.rand(100.0)
  let ret = tree.nearestNeighbours([x, y], n)
  for (pt, value, dist) in ret:
    echo fmt"point={pt}, value={value}, dist={dist}"

# Perform withinRadius searches
var ret2 = tree.withinRadius([point.x, point.y], radius=5.0, sortResults=true)
for (pt, value, dist) in ret2:
  echo fmt"point={pt}, value={value}, dist={dist}"

# Perform withinRange searches
var
  min: array[2, float] = [0.0, 0.0]
  max: array[2, float] = [10.0, 10.0]
  hyperRect = newHyperRectangle(min, max)

var ret = tree.withinRange(hyperRect)
for (pt, value) in ret:
  echo fmt"point={pt}, value={value}"

Types

KdPoint = array[K, float]: A KdPoint is a location in K-dimensional space.
DistFunc = proc (x, y: KdPoint): float {...}{.closure.}: A distance function used to calculate the distance between two KdPoints, returning a float
KdTree[T] = object root*: KdNode[T] len: int distFunc: DistFunc: A k-d tree data structure that allows efficient spatial querying on point distributions. The Current implementation is designed for 2-D point data, although other dimensionality is possible simply by modifying the const K.
HyperRectangle = object min*: KdPoint max*: KdPoint: A HyperRectangle is used by the withinRange search function to identify multi-deminsional ranges.

Consts

K = 2: K is the dimensionality of the points in this package's K-D trees.

Procs

proc nearestNeighbour[T](tree: var KdTree[T]; point: KdPoint): (KdPoint, T, float)

Returns the nearest neighbour of an input target point, the data associated with the nearest neighbour, and the distance between the target point and the nearest neighbour. Notice that the returned distance uses the distance metric based on the distFunc parameter when the tree is created. By default, and if unspecified, this metric is the squared distance.

let x = 100.0
let y = 25.0
let (pt, values, dist) = tree.nearestNeighbour([x, y])

proc nearestNeighbours[T](tree: var KdTree[T]; point: KdPoint; numNeighbours: int): seq[ (KdPoint, T, float)]

Returns a specified number (numNeighbours) of nearest neighbours of a target point (point). Each return point is accompanied by the associated data, and the distance between the target and return points. Notice that the returned distance uses the distance metric based on the distFunc parameter when the tree is created. By default, and if unspecified, this metric is the squared distance.

let x = 100.0
let y = 25.0
let ret = tree.nearestNeighbours([x, y], numNeighbours=5)
for (pt, value, dist) in ret:
  echo fmt"point={pt}, value={values}, dist={dist}"

proc withinRadius[T](tree: var KdTree[T]; point: KdPoint; radius: float; sortResults = false): seq[(KdPoint, T, float)]

Returns all of the points contained in the tree that are within a specified radius of a target point. By default, the returned points are in an arbitrary order, unless the sortResults parameter is set to true, in which case the return points will be sorted from nearest to farthest from the target. Notice that the returned distance uses the distance metric based on the distFunc parameter when the tree is created. By default, and if unspecified, this metric is the squared distance. Importantly, the radius parameter must be a distance measured using the same distance metric as specified by the distFunc parameter when the tree is created.

let x = 100.0
let y = 25.0
var ret = tree.withinRadius([x, y], radius=5.0, sortResults=true)
for (pt, value, dist) in ret:
   echo fmt"point={pt}, value={value}, dist={dist}"

Funcs

func newKdTree[T](pointData: openArray[(KdPoint, T)]; distFunc: DistFunc = sqrDist): KdTree[ T]

Constructs a k-d tree by bulk-loading an array of point-data tuples, where the associated data is of any generic type T. Notice that this way of constructing a KdTree should be preferred over adding points individually because the resulting tree will be balanced, which will make for more efficient search operations. The default distFunc is the squared distance, which is returned from each search function.

let pointsAndValues = [([2.0, 3.0], 1),
                      ([5.0, 4.0], 2),
                      ([9.0, 6.0], 3),
                      ([4.0, 7.0], 4),
                      ([8.0, 1.0], 5),
                      ([7.0, 2.0], 6)]

var tree = newKdTree[int](pointsAndValues)

# A custom distance function; the default is a squared distance
proc myDistFunc(self, other: array[2, float]): float =
  result = 0.0
  for i in 0..<len(self):
    result += (self[i] - other[i]) * (self[i] - other[i])
  result = sqrt(result)

var tree = newKdTree[int](pointsAndValues, distFunc=myDistFunc)

func newKdTree[T](points: openArray[KdPoint]; data: openArray[T]; distFunc: DistFunc = sqrDist): KdTree[T]

Constructs a k-d tree by bulk-loading arrays of points and associated data values of any generic type T. Notice that this way of constructing a KdTree should be preferred over adding points individually because the resulting tree will be balanced, which will make for more efficient search operations. The default distFunc is the squared distance, which is returned from each search function.

let points = [[2.0, 3.0], [5.0, 4.0], [9.0, 6.0], [4.0, 7.0], [8.0, 1.0], [7.0, 2.0]]
let values = [1, 2, 3, 4, 5, 6]

var tree = newKdTree[int](points, values)

# A custom distance function; the default is a squared distance
proc myDistFunc(self, other: array[2, float]): float =
  result = 0.0
  for i in 0..<len(self):
    result += (self[i] - other[i]) * (self[i] - other[i])
  result = sqrt(result)

var tree = newKdTree[int](points, values, distFunc=myDistFunc)

func add[T](tree: var KdTree[T]; point: KdPoint; data: T)

This function can be used to add single points, and their associated data of type T to an existing KdTree object. Notice that this can result in an unbalanced tree which is suboptimal for search operation efficiency.

func len[T](tree: KdTree[T]): int

Returns the number of nodes contained within the KdTree.

func height[T](tree: var KdTree[T]): int

Returns the height of the KdTree.

func isBalanced[T](tree: var KdTree[T]): int

Returns the value of the left tree height - right tree height. The larger the value magnitude, the more unbalanced the tree is (some say an unbalanced tree is any with an absolute magnitude greater than 1). The sign indicates the direction of skew, with negative values indicating a left-skewed tree and positive values indicated a right-skewed tree.

func rebalance[T](tree: var KdTree[T])

Re-balances an unbalanced KdTree. Note that the original tree structure can be completely modified by this function. Use this function after adding a significant number of individual nodes to the tree with the add function.

func newHyperRectangle(min: KdPoint; max: KdPoint): HyperRectangle {...}{.raises: [], tags: [].}

Creates a new HyperRectangle

func withinRange[T](tree: var KdTree[T]; rectangle: HyperRectangle): seq[(KdPoint, T)]

Returns all of the points contained in the tree that are within a target HyperRectangle.

var
  min: array[2, float] = [0.0, 0.0]
  max: array[2, float] = [100.0, 100.0]
  hyperRect = newHyperRectangle(min, max)

var ret = tree.withinRange(hyperRect)
for (pt, value) in ret:
  echo fmt"point={pt}, value={value}"