go-kdtree's Introduction

KDTree

Golang implementation of KD tree (https://en.wikipedia.org/wiki/K-d_tree) data structure

Getting started

Use go tool to install the package in your packages tree:

go get github.com/hongshibao/go-kdtree

Then you can use it in import section of your Go programs:

import "github.com/hongshibao/go-kdtree"

The package name is kdtree.

Basic example

First you need to implement the Point interface:

type Point interface {
	// Return the total number of dimensions
	Dim() int
	// Return the value X_{dim}, dim is started from 0
	GetValue(dim int) float64
	// Return the distance between two points
	Distance(p Point) float64
	// Return the distance between the point and the plane X_{dim}=val
	PlaneDistance(val float64, dim int) float64
}

Here is an example of implementing Point interface with square of Euclidean distance as the Distance definition:

type EuclideanPoint struct {
	Point
	Vec []float64
}

func (p *EuclideanPoint) Dim() int {
	return len(p.Vec)
}

func (p *EuclideanPoint) GetValue(dim int) float64 {
	return p.Vec[dim]
}

func (p *EuclideanPoint) Distance(other Point) float64 {
	var ret float64
	for i := 0; i < p.Dim(); i++ {
		tmp := p.GetValue(i) - other.GetValue(i)
		ret += tmp * tmp
	}
	return ret
}

func (p *EuclideanPoint) PlaneDistance(val float64, dim int) float64 {
	tmp := p.GetValue(dim) - val
	return tmp * tmp
}

Now you can create KD-tree from a list of points and get a list of k nearest neighbours for a target point:

func NewEuclideanPoint(vals ...float64) *EuclideanPoint {
	ret := &EuclideanPoint{}
	for _, val := range vals {
		ret.Vec = append(ret.Vec, val)
	}
	return ret
}

func main() {
	p1 := NewEuclideanPoint(0.0, 0.0, 0.0)
	p2 := NewEuclideanPoint(0.0, 0.0, 1.0)
	p3 := NewEuclideanPoint(0.0, 1.0, 0.0)
	p4 := NewEuclideanPoint(1.0, 0.0, 0.0)
	points := make([]Point, 0)
	points = append(points, p1)
	points = append(points, p2)
	points = append(points, p3)
	points = append(points, p4)
	tree := NewKDTree(points)
	targetPoint := NewEuclideanPoint(0.0, 0.0, 0.1)
	neighbours := tree.KNN(targetPoint, 2)
	for idx, p := range neighbours {
		fmt.Printf("Point %d: (%f", idx, p.GetValue(0))
		for i := 1; i < p.Dim(); i++ {
			fmt.Printf(", %f", p.GetValue(i))
		}
		fmt.Println(")")
	}
}

The returned k nearest neighbours are sorted by their distance with the target point.

go-kdtree's People

Contributors

Stargazers

Watchers

go-kdtree's Issues

Provide a working example

Could you please provide a working example to get novice users (like me)?

KNN returns points number less than specified

Hi, we've tried to use your library but faced the problem.

I've added one more point to your test case and it returns the wrong response:

func TestKNN(t *testing.T) {
	points := []Point{
		NewEuclideanPoint(0.0, 0.0, 0.0),
		NewEuclideanPoint(0.0, 0.0, 1.0),
		NewEuclideanPoint(0.0, 1.0, 0.0),
		NewEuclideanPoint(1.0, 0.0, 0.0),
		NewEuclideanPoint(0.0, 0.0, 0.0),
		NewEuclideanPoint(0.0, 0.0, 0.1),
		NewEuclideanPoint(1.0, 1.0, 1.0),
		NewEuclideanPoint(0.1, 0.1, 0.1), // aditional point
	}

	tree := NewKDTree(points)

	ans := tree.KNN(NewEuclideanPoint(0.0, 0.0, 0.0), 7)
	if len(ans) != 7 {
		t.Errorf("expected 7 points, actual: %v", len(ans))
	}
}