kdtree

README

kdtree

A k-d tree implementation in Go with:

• n-dimensional points
• k-nearest neighbor search
• range search
• remove without rebuilding the whole subtree
• data attached to the points
• using own structs by implementing a simple 2 function interface

Usage

go get github.com/kyroy/kdtree

import "github.com/kyroy/kdtree"

Implement the kdtree.Point interface

// Point specifies one element of the k-d tree.
type Point interface {
	// Dimensions returns the total number of dimensions
	Dimensions() int
	// Dimension returns the value of the i-th dimension
	Dimension(i int) float64
}

points.Point2d

type Data struct {
	value string
}

func main() {
	tree := kdtree.New([]kdtree.Point{
		&points.Point2D{X: 3, Y: 1},
		&points.Point2D{X: 5, Y: 0},
		&points.Point2D{X: 8, Y: 3},
	})

	// Insert
	tree.Insert(&points.Point2D{X: 1, Y: 8})
	tree.Insert(&points.Point2D{X: 7, Y: 5})

	// KNN (k-nearest neighbor)
	fmt.Println(tree.KNN(&points.Point{Coordinates: []float64{1, 1, 1}}, 2))
	// [{3.00 1.00} {5.00 0.00}]
	
	// RangeSearch
	fmt.Println(tree.RangeSearch(kdrange.New(1, 8, 0, 2)))
	// [{5.00 0.00} {3.00 1.00}]
    
	// Points
	fmt.Println(tree.Points())
	// [{3.00 1.00} {1.00 8.00} {5.00 0.00} {8.00 3.00} {7.00 5.00}]

	// Remove
	fmt.Println(tree.Remove(&points.Point2D{X: 5, Y: 0}))
	// {5.00 0.00}

	// String
	fmt.Println(tree)
	// [[{1.00 8.00} {3.00 1.00} [<nil> {8.00 3.00} {7.00 5.00}]]]

	// Balance
	tree.Balance()
	fmt.Println(tree)
	// [[[{3.00 1.00} {1.00 8.00} <nil>] {7.00 5.00} {8.00 3.00}]]
}

n-dimensional Points (points.Point)

type Data struct {
	value string
}

func main() {
    tree := kdtree.New([]kdtree.Point{
        points.NewPoint([]float64{7, 2, 3}, Data{value: "first"}),
        points.NewPoint([]float64{3, 7, 10}, Data{value: "second"}),
        points.NewPoint([]float64{4, 6, 1}, Data{value: "third"}),
    })
    
    // Insert
    tree.Insert(points.NewPoint([]float64{12, 4, 6}, Data{value: "fourth"}))
    tree.Insert(points.NewPoint([]float64{8, 1, 0}, Data{value: "fifth"}))
    
    // KNN (k-nearest neighbor)
    fmt.Println(tree.KNN(&points.Point{Coordinates: []float64{1, 1, 1}}, 2))
    // [{[4 6 1] {third}} {[7 2 3] {first}}]
    
    // RangeSearch
    fmt.Println(tree.RangeSearch(kdrange.New(1, 15, 1, 5, 0, 5)))
    // [{[7 2 3] {first}} {[8 1 0] {fifth}}]
    
    // Points
    fmt.Println(tree.Points())
    // [{[3 7 10] {second}} {[4 6 1] {third}} {[8 1 0] {fifth}} {[7 2 3] {first}} {[12 4 6] {fourth}}]

    // Remove
    fmt.Println(tree.Remove(points.NewPoint([]float64{3, 7, 10}, nil)))
    // {[3 7 10] {second}}

    // String
    fmt.Println(tree)
    // [[<nil> {[4 6 1] {third}} [{[8 1 0] {fifth}} {[7 2 3] {first}} {[12 4 6] {fourth}}]]]

    // Balance
    tree.Balance()
    fmt.Println(tree)
    // [[[{[7 2 3] {first}} {[4 6 1] {third}} <nil>] {[8 1 0] {fifth}} {[12 4 6] {fourth}}]]
}

Documentation

Overview

Package kdtree implements a k-d tree data structure.

Index

• type KDTree
  • func New(points []Point) *KDTree
  • func (t *KDTree) Balance()
  • func (t *KDTree) Insert(p Point)
  • func (t *KDTree) KNN(p Point, k int) []Point
  • func (t *KDTree) Points() []Point
  • func (t *KDTree) RangeSearch(r kdrange.Range) []Point
  • func (t *KDTree) Remove(p Point) Point
  • func (t *KDTree) String() string
• type Point

Types

type KDTree

type KDTree struct {
	// contains filtered or unexported fields
}

KDTree represents the k-d tree.

func New

func New(points []Point) *KDTree

New returns a balanced k-d tree.

func (*KDTree) Balance

func (t *KDTree) Balance()

Balance rebalances the k-d tree by recreating it.

func (*KDTree) Insert

func (t *KDTree) Insert(p Point)

Insert adds a point to the k-d tree.

func (*KDTree) KNN

func (t *KDTree) KNN(p Point, k int) []Point

KNN returns the k-nearest neighbours of the given point. The points are sorted by the distance to the given points. Starting with the nearest.

func (*KDTree) Points

func (t *KDTree) Points() []Point

Points returns all points in the k-d tree. The tree is traversed in-order.

func (*KDTree) RangeSearch

func (t *KDTree) RangeSearch(r kdrange.Range) []Point

RangeSearch returns all points in the given range r.

Returns an empty slice when input is nil or len(r) does not equal Point.Dimensions().

func (*KDTree) Remove

func (t *KDTree) Remove(p Point) Point

Remove removes and returns the first point from the tree that equals the given point p in all dimensions. Returns nil if not found.

func (*KDTree) String

func (t *KDTree) String() string

String returns a string representation of the k-d tree.

type Point

type Point interface {
	// Dimensions returns the total number of dimensions.
	Dimensions() int
	// Dimension returns the value of the i-th dimension.
	Dimension(i int) float64
}

Point specifies one element of the k-d tree.