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/kdtreeimport "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}}]]
}