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

• func EuclideanDistance(p1, p2 Point) float64
• type DistanceFunc
• type KDTree
  • func New(points []Point) *KDTree
  • func NewCustom(points []Point, distance DistanceFunc) *KDTree
  • func (t *KDTree) Balance()
  • func (t *KDTree) Insert(p Point)
  • func (t *KDTree) KNN(p Point, k int) []Point
  • func (t *KDTree) LineTrace(start Point, end Point, radius float64) (*Point, float64)
  • 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

Functions

func EuclideanDistance

func EuclideanDistance(p1, p2 Point) float64

Types

type DistanceFunc

type DistanceFunc func(p1, p2 Point) float64

Used for different distance functions (manhattan, euclidean, etc.)

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 NewCustom

func NewCustom(points []Point, distance DistanceFunc) *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) LineTrace

func (t *KDTree) LineTrace(start Point, end Point, radius float64) (*Point, float64)

RayTrace returns the first point (with the radius) that intersects this LineSegment from start to end and the dist-factor from [0,1], or "max float" if no hit

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.