kyroy / kdtree Goto Github PK

package main

import (
	"fmt"
	"math"
	"math/rand"
	"time"

	"github.com/kyroy/kdtree"
	"github.com/kyroy/kdtree/points"
)

func distance(a, b *points.Point2D) float64 {
	return math.Sqrt(math.Pow(a.X - b.X, 2) + math.Pow(a.Y - b.Y, 2))
}

func dumpNNs(p kdtree.Point, nns []kdtree.Point) {
	p2d1 := p.(*points.Point2D)
	fmt.Printf("KNN((%3.2f, %3.2f), %d)\n", p2d1.X, p2d1.Y, len(nns))

	for i, n := range nns {
		p2d2 := n.(*points.Point2D)
		fmt.Printf("   NN %d: distance((%3.2f, %3.2f)->(%3.2f, %3.2f)) = %3.2f\n",
			i, p2d1.X, p2d1.Y, p2d2.X, p2d2.Y, distance(p2d1, p2d2))
	}
}

func main() {
	w, h := 300, 300
	maxSize := 10
	maxTime := time.Minute

	t := kdtree.New([]kdtree.Point{})
	arr := make([]kdtree.Point, 0, maxSize)

	many := maxSize
	few := maxSize - 1

	for ; few > 0; few-- {
		fmt.Printf("Getting %d NNs out of %d entries\n", few, maxSize)
		start := time.Now()
		for i := 0; ; i++ {
			p := &points.Point2D{
				X: float64(rand.Intn(w / 2) + w / 2),
				Y: float64(rand.Intn(h / 2) + h / 2),
			}

			// Two KNN queries
			fewNN := t.KNN(p, few)
			manyNN := t.KNN(p, many)

			// Check if the nearest is the same
			if len(fewNN) > 0 {
				if distance(p, fewNN[0].(*points.Point2D)) >
					distance(p, manyNN[0].(*points.Point2D)) {
					fmt.Println("Failure on iteration", i)
					dumpNNs(p, fewNN)
					dumpNNs(p, manyNN)

					// Balance the tree and try again
					t.Balance()
					fmt.Println("After balance:")
					fewNN = t.KNN(p, few)
					manyNN = t.KNN(p, many)
					dumpNNs(p, fewNN)
					dumpNNs(p, manyNN)

					break
				}
			}

			// Add in the new point
			arr = append(arr, p)
			//fmt.Printf("Insert (%3.2f, %3.2f)\n", p.X, p.Y)
			t.Insert(p)

			// Limit the max number of elements - which will also
			// introduce some churn in the tree
			if len(arr) > maxSize {
				idx := rand.Intn(len(arr))
				//fmt.Printf("Remove (%3.2f, %3.2f)\n", p.X, p.Y)
				t.Remove(arr[idx])
				p = arr[idx].(*points.Point2D)
				arr[idx] = arr[len(arr)-1]
				arr = arr[:len(arr)-1]
			}

			if since := time.Since(start); since > maxTime {
				fmt.Printf("No failure after %d iterations (%v)\n", i, since)
				break
			}
		}
		fmt.Println("================================================")
		t.Balance()
	}
}