diff --git a/go.mod b/go.mod index 4840da8..d6dd020 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.18 require ( github.com/downflux/go-geometry v0.10.2 + github.com/downflux/go-pq v0.1.1 github.com/google/go-cmp v0.5.8 github.com/kyroy/kdtree v0.0.0-20200419114247-70830f883f1d ) diff --git a/go.sum b/go.sum index 912433d..96a1da9 100644 --- a/go.sum +++ b/go.sum @@ -1,15 +1,22 @@ +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/downflux/go-geometry v0.10.2 h1:Z79Khzl6AKMSMLnM5xG75fEOL1fmIWlF14+8j+r01D0= github.com/downflux/go-geometry v0.10.2/go.mod h1:XWTzSaMiRMAxupAR+cXAsa1Q75TCSp1Shc/ydsJ0xVE= +github.com/downflux/go-pq v0.1.1 h1:dwWLUMBiUTrZO+TlD+oHPQqsqQDo3f7Yyo8nvO1CV7Y= +github.com/downflux/go-pq v0.1.1/go.mod h1:6amZrMlYUeY1Ba9T+0A11y2BBENzsUqMmu4o5CMndMk= github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg= github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/jupp0r/go-priority-queue v0.0.0-20160601094913-ab1073853bde h1:+5PMaaQtDUwOcJIUlmX89P0J3iwTvErTmyn5WghzXAQ= github.com/jupp0r/go-priority-queue v0.0.0-20160601094913-ab1073853bde/go.mod h1:RDgD/dfPmIwFH0qdUOjw71HjtWg56CtyLIoHL+R1wJw= github.com/kyroy/kdtree v0.0.0-20200419114247-70830f883f1d h1:1n5M/49q9H6QtNJiiVL/W5mqgT1UdlGQ7oLP+DkJ1vs= github.com/kyroy/kdtree v0.0.0-20200419114247-70830f883f1d/go.mod h1:6oJGQK7VSg3RxSQ7QspgqpCmKjIbAslgT2wBXbFJUZw= github.com/kyroy/priority-queue v0.0.0-20180327160706-6e21825e7e0c h1:1c7+XOOGQ19cXjZ1Ss/irljQxgPvb+8z+jNEprCXl20= github.com/kyroy/priority-queue v0.0.0-20180327160706-6e21825e7e0c/go.mod h1:R477L6j2/dUcE0q0aftk0kR5Xt93W7g1066AodcJhEo= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= diff --git a/internal/knn/knn.go b/internal/knn/knn.go index bef6811..7cb8f11 100644 --- a/internal/knn/knn.go +++ b/internal/knn/knn.go @@ -5,7 +5,7 @@ import ( "github.com/downflux/go-kd/internal/node" "github.com/downflux/go-kd/point" - "github.com/downflux/go-kd/point/pq" + "github.com/downflux/go-pq/pq" "github.com/downflux/go-kd/vector" vnd "github.com/downflux/go-geometry/nd/vector" diff --git a/point/pq/pq.go b/point/pq/pq.go deleted file mode 100644 index 7dff0af..0000000 --- a/point/pq/pq.go +++ /dev/null @@ -1,94 +0,0 @@ -package pq - -import ( - "container/heap" - "math" - - "github.com/downflux/go-kd/point" -) - -type node[T point.P] struct { - p T - index int - weight float64 -} - -// max implements the heap.Interface. -// -// See https://pkg.go.dev/container/heap for the heap implementation from which -// this was shamelessly stolen. -type max[T point.P] []*node[T] - -func (heap *max[T]) Len() int { return len(*heap) } -func (heap *max[T]) Less(i, j int) bool { return (*heap)[i].weight > (*heap)[j].weight } -func (heap *max[T]) Swap(i, j int) { - (*heap)[i], (*heap)[j] = (*heap)[j], (*heap)[i] - (*heap)[i].index = i - (*heap)[j].index = j -} -func (heap *max[T]) Push(x any) { - m := x.(*node[T]) - m.index = heap.Len() - *heap = append(*heap, m) -} -func (heap *max[T]) Pop() any { - h := *heap - n := len(h) - m := h[n-1] - - h[n-1] = nil - m.index = -1 - - *heap = h[0 : n-1] - return m -} - -// PQ is an exported priority queue of k-D tree nodes with a maximum queue size. -// PQ tracks a specified number points with the smallest given priorities -- -// attempting to add a point with a larger priority will result in an effective -// no-op. -type PQ[T point.P] struct { - heap *max[T] - n int -} - -func New[T point.P](size int) *PQ[T] { - h := max[T](make([]*node[T], 0, size)) - pq := &PQ[T]{ - heap: &h, - n: size, - } - heap.Init(pq.heap) - return pq -} - -func (pq *PQ[T]) Len() int { return pq.heap.Len() } -func (pq *PQ[T]) Empty() bool { return pq.Len() == 0 } -func (pq *PQ[T]) Full() bool { return pq.Len() >= pq.n } - -// Priority calculates the current highest priority of the queue. -func (pq *PQ[T]) Priority() float64 { - if pq.Empty() { - return math.Inf(0) - } - - // See https://groups.google.com/g/golang-nuts/c/sy1p8SfyPoY. - return (*pq.heap)[0].weight -} - -// Push adds a new point into the queue. -// -// The queue will enforce the struct size constraint by removing elements frmo -// itself until the constraint is satisfied. -func (pq *PQ[T]) Push(p T, priority float64) { - heap.Push(pq.heap, &node[T]{ - p: p, - weight: priority, - }) - for !pq.Empty() && pq.Len() > pq.n { - pq.Pop() - } -} - -// Pop removes the node with the highest priority from the queue. -func (pq *PQ[T]) Pop() T { return heap.Pop(pq.heap).(*node[T]).p } diff --git a/point/pq/pq_test.go b/point/pq/pq_test.go deleted file mode 100644 index 85e047a..0000000 --- a/point/pq/pq_test.go +++ /dev/null @@ -1,184 +0,0 @@ -package pq - -import ( - "container/heap" - "testing" - - "github.com/downflux/go-kd/point/mock" - "github.com/google/go-cmp/cmp" -) - -var _ heap.Interface = &max[mock.P]{} - -func TestPQ(t *testing.T) { - type item struct { - p mock.P - priority float64 - } - - type config struct { - name string - data []item - size int - want []mock.P - } - - configs := []config{ - { - name: "Null", - data: nil, - want: nil, - }, - - { - name: "Trivial", - data: []item{ - { - p: mock.P{ - X: mock.U(1), - }, - priority: 0, - }, - }, - size: 1, - want: []mock.P{ - mock.P{ - X: mock.U(1), - }, - }, - }, - { - name: "Trivial/NoSize", - data: []item{ - { - p: mock.P{ - X: mock.U(1), - }, - priority: 0, - }, - }, - size: 0, - want: nil, - }, - - { - name: "Sorted", - data: []item{ - { - p: mock.P{X: mock.U(0)}, - priority: 5, - }, - { - p: mock.P{X: mock.U(-1)}, - priority: 4, - }, - { - p: mock.P{X: mock.U(-2)}, - priority: 3, - }, - { - p: mock.P{X: mock.U(-3)}, - priority: 2, - }, - { - p: mock.P{X: mock.U(-4)}, - priority: 1, - }, - }, - size: 5, - want: []mock.P{ - mock.P{X: mock.U(0)}, - mock.P{X: mock.U(-1)}, - mock.P{X: mock.U(-2)}, - mock.P{X: mock.U(-3)}, - mock.P{X: mock.U(-4)}, - }, - }, - { - name: "Sorted/Reverse", - data: []item{ - { - p: mock.P{X: mock.U(0)}, - priority: 1, - }, - { - p: mock.P{X: mock.U(-1)}, - priority: 2, - }, - { - p: mock.P{X: mock.U(-2)}, - priority: 3, - }, - { - p: mock.P{X: mock.U(-3)}, - priority: 4, - }, - { - p: mock.P{X: mock.U(-4)}, - priority: 5, - }, - }, - size: 5, - want: []mock.P{ - mock.P{X: mock.U(-4)}, - mock.P{X: mock.U(-3)}, - mock.P{X: mock.U(-2)}, - mock.P{X: mock.U(-1)}, - mock.P{X: mock.U(0)}, - }, - }, - { - name: "Sorted/Shuffled", - data: []item{ - { - p: mock.P{X: mock.U(0)}, - priority: 4, - }, - { - p: mock.P{X: mock.U(-1)}, - priority: 2, - }, - { - p: mock.P{X: mock.U(-2)}, - priority: 5, - }, - { - p: mock.P{X: mock.U(-3)}, - priority: 1, - }, - { - p: mock.P{X: mock.U(-4)}, - priority: 3, - }, - }, - size: 5, - want: []mock.P{ - mock.P{X: mock.U(-2)}, - mock.P{X: mock.U(0)}, - mock.P{X: mock.U(-4)}, - mock.P{X: mock.U(-1)}, - mock.P{X: mock.U(-3)}, - }, - }, - } - - for _, c := range configs { - t.Run(c.name, func(t *testing.T) { - pq := New[mock.P](c.size) - for _, d := range c.data { - pq.Push(d.p, d.priority) - } - - var got []mock.P - for !pq.Empty() { - got = append(got, pq.Pop()) - } - - if diff := cmp.Diff( - c.want, got, - ); diff != "" { - t.Errorf("Pop() mismatch (-want +got):\n%v", diff) - } - }) - } -}