Updated geoindex

This commit is contained in:
tidwall 2019-09-19 11:30:49 -07:00
parent 639f6e2deb
commit 5abadd72a3
3 changed files with 85 additions and 26 deletions

6
Gopkg.lock generated
View File

@ -213,15 +213,15 @@
version = "v1.1.0" version = "v1.1.0"
[[projects]] [[projects]]
digest = "1:4305de55e110aa13ec68a246b12e512041fe92440e78066fcea93ecf2a68320b" digest = "1:eab1a01c55a3428f83e16e92f902ffbeae143e19e080a6a1117532f7908f7579"
name = "github.com/tidwall/geoindex" name = "github.com/tidwall/geoindex"
packages = [ packages = [
".", ".",
"child", "child",
] ]
pruneopts = "" pruneopts = ""
revision = "e56705dcd2788d8eb431e8cb15295bfd0a298976" revision = "6fc1984907cad925af47d09fdc0cadc70f875cfe"
version = "v1.0.1" version = "v1.1.0"
[[projects]] [[projects]]
digest = "1:ddb305f09be3613fd1bf9fd8d6d0713f2fd28b5af596437b3d7de2366bbee870" digest = "1:ddb305f09be3613fd1bf9fd8d6d0713f2fd28b5af596437b3d7de2366bbee870"

View File

@ -73,31 +73,50 @@ func (index *Index) Children(parent interface{}, reuse []child.Child) (
return index.tree.Children(parent, reuse) return index.tree.Children(parent, reuse)
} }
// Nearby performs a kNN-type operation on the index. It's expected that the // Nearby performs a kNN-type operation on the index.
// caller provides the `dist` function, which is used to calculate the // It's expected that the caller provides its own the `algo` function, which
// distance from a node or item to another object. The other object is unknown // is used to calculate a distance to data. The `add` function should be
// this operation, but is expected to be known by the caller. The iter will // called by the caller to "return" the data item along with a distance.
// return all items from the smallest dist to the largest dist. // The `iter` function will return all items from the smallest dist to the
// largest dist.
// Take a look at the SimpleBoxAlgo function for a usage example.
func (index *Index) Nearby( func (index *Index) Nearby(
algo func(min, max [2]float64, data interface{}, item bool) (dist float64), algo func(
min, max [2]float64, data interface{}, item bool,
add func(min, max [2]float64, data interface{}, item bool, dist float64),
),
iter func(min, max [2]float64, data interface{}, dist float64) bool, iter func(min, max [2]float64, data interface{}, dist float64) bool,
) { ) {
var q queue var q queue
var parent interface{} var parent interface{}
var children []child.Child var children []child.Child
var added []qnode
add := func(min, max [2]float64, data interface{}, item bool, dist float64) {
added = append(added, qnode{
dist: dist,
child: child.Child{
Data: data,
Min: min,
Max: max,
Item: item,
},
})
}
for { for {
// gather all children for parent // gather all children for parent
children = index.tree.Children(parent, children[:0]) children = index.tree.Children(parent, children[:0])
for _, child := range children { for _, child := range children {
q.push(qnode{ added = added[:0]
dist: algo(child.Min, child.Max, child.Data, child.Item), algo(child.Min, child.Max, child.Data, child.Item, add)
child: child, for _, node := range added {
filled: true, q.push(node)
}) }
} }
for { for {
node := q.pop() node, ok := q.pop()
if !node.filled { if !ok {
// nothing left in queue // nothing left in queue
return return
} }
@ -128,9 +147,8 @@ func (index *Index) Bounds() (min, max [2]float64) {
// Priority Queue ordered by dist (smallest to largest) // Priority Queue ordered by dist (smallest to largest)
type qnode struct { type qnode struct {
dist float64 dist float64
child child.Child child child.Child
filled bool
} }
type queue struct { type queue struct {
@ -156,9 +174,9 @@ func (q *queue) push(node qnode) {
q.len++ q.len++
} }
func (q *queue) pop() qnode { func (q *queue) pop() (qnode, bool) {
if q.len == 0 { if q.len == 0 {
return qnode{} return qnode{}, false
} }
n := q.nodes[1] n := q.nodes[1]
q.nodes[1] = q.nodes[q.len] q.nodes[1] = q.nodes[q.len]
@ -177,7 +195,7 @@ func (q *queue) pop() qnode {
q.nodes[i] = q.nodes[k] q.nodes[i] = q.nodes[k]
i = k i = k
} }
return n return n, true
} }
// Scan iterates through all data in tree in no specified order. // Scan iterates through all data in tree in no specified order.
@ -187,12 +205,18 @@ func (index *Index) Scan(
index.tree.Scan(iter) index.tree.Scan(iter)
} }
// SimpleBoxAlgo ... // SimpleBoxAlgo performs box-distance algorithm on rectangles.
func SimpleBoxAlgo(targetMin, targetMax [2]float64) ( func SimpleBoxAlgo(targetMin, targetMax [2]float64) (
dist func(min, max [2]float64, data interface{}, item bool) (dist float64), algo func(
min, max [2]float64, data interface{}, item bool,
add func(min, max [2]float64, data interface{}, item bool, dist float64),
),
) { ) {
return func(min, max [2]float64, data interface{}, item bool) float64 { return func(
return boxDist(targetMin, targetMax, min, max) min, max [2]float64, data interface{}, item bool,
add func(min, max [2]float64, data interface{}, item bool, dist float64),
) {
add(min, max, data, item, boxDist(targetMin, targetMax, min, max))
} }
} }

View File

@ -314,6 +314,41 @@ func testBoxesVarious(t *testing.T, boxes []tBox, label string) {
ldist = dist ldist = dist
} }
// test bounds
min := boxes3[0].min
max := boxes3[0].max
for _, box := range boxes3 {
if box.min[0] < min[0] {
min[0] = box.min[0]
}
if box.min[1] < min[1] {
min[1] = box.min[1]
}
if box.max[0] > max[0] {
max[0] = box.max[0]
}
if box.max[1] > max[1] {
max[1] = box.max[1]
}
}
min2, max2 := tr.Bounds()
if min2 != min || max2 != max {
t.Fatalf("expected %v,%v, got %v,%v", min, max, min2, max2)
}
// test nearby, but stop after one
var one tBox
tr.Nearby(
SimpleBoxAlgo(centerMin, centerMax),
func(min, max [2]float64, value interface{}, dist float64) bool {
one = value.(tBox)
return false
},
)
if one != boxes3[0] {
t.Fatalf("expected %v, got %v", boxes[0], one)
}
} }
func TestRandomBoxes(t *testing.T) { func TestRandomBoxes(t *testing.T) {