Let's say you have a large set of 3D points called data points, and you'd like to be able to quickly perform point queries on the data points. One example of a point query is the nearest neighbor query: given a set of data points points and a query point p, which point in points is closest to p?

We can efficiently solve the nearest neighbor query (along with many other types of point queries) if we appropriately organize the data points. One such method of organization is called the k-d tree algorithm, which is implemented in this module.

Let's say you have a list of 3D data points, and each point is of type Point3d:

data Point3d = Point3d { x :: Double
                       , y :: Double
                       , z :: Double
                       } deriving Show

We call a point's individual values axis values (i.e., x, y, and z in the case of Point3d).

In order to generate a k-d tree of Point3d's, we need to define a PointAsListFn that expresses the point's axis values as a list:

point3dAsList :: Point3d -> [Double]
point3dAsList (Point3d x y z) = [x, y, z]

Now we can build a KdTree structure from a list of data points and perform a nearest neighbor query as follows:

>>> let dataPoints = [(Point3d 0.0 0.0 0.0), (Point3d 1.0 1.0 1.0)]

>>> let kdt = buildKdTree point3dAsList dataPoints

>>> let queryPoint = Point3d 0.1 0.1 0.1

>>> nearestNeighbor kdt queryPoint
Point3d {x = 0.0, y = 0.0, z = 0.0}

The KdTree structure is meant for static sets of data points. If you need to insert points into an existing k-d tree, check out Data.KdTree.Dynamic.KdTree.

If you need to associate additional data with each point in the tree (i.e., points are keys associated with values), check out Data.KdMap.Static.KdMap and Data.KdMap.Dynamic.KdMap for static and dynamic variants of this functionality. Please do not try to fake this functionality with a KdTree by augmenting your point type with the extra data; you're gonna have a bad time.

You may have noticed in the previous use case that we never specified what "nearest" means for our points. By default, buildKdTree uses a Euclidean distance function that is sufficient in most cases. However, point queries are typically faster on a KdTree built with a user-specified custom distance function. Let's generate a KdTree using a custom distance function.

One idiosyncrasy about KdTree is that custom distance functions are actually specified as squared distance functions (SquaredDistanceFn). This means that your custom distance function must return the square of the actual distance between two points. This is for efficiency: regular distance functions often require expensive square root computations, whereas in our case, the squared distance works fine and doesn't require computing any square roots. Here's an example of a squared distance function for Point3d:

point3dSquaredDistance :: Point3d -> Point3d -> Double
point3dSquaredDistance (Point3d x1 y1 z1) (Point3d x2 y2 z2) =
  let dx = x1 - x2
      dy = y1 - y2
      dz = z1 - z2
  in  dx * dx + dy * dy + dz * dz

We can build a KdTree using our custom distance function as follows:

>>> let kdt = buildKdTreeWithDistFn point3dAsList point3dSquaredDistance points

In the above examples, we used a point type with axis values of type Double. We can in fact use axis values of any type that is an instance of the Real typeclass. This means you can use points that are composed of Doubles, Ints, Floats, and so on:

data Point2i = Point2i Int Int

point2iAsList :: Point2i -> [Int]
point2iAsList (Point2i x y) = [x, y]

kdt :: [Point2i] -> KdTree Int Point2i
kdt dataPoints = buildKdTree point2iAsList dataPoints