package geoindex import ( "fmt" "github.com/tidwall/geoindex/child" ) // Interface is a tree-like structure that contains geospatial data type Interface interface { // Insert an item into the structure Insert(min, max [2]float64, data interface{}) // Delete an item from the structure Delete(min, max [2]float64, data interface{}) // Search the structure for items that intersects the rect param Search( min, max [2]float64, iter func(min, max [2]float64, data interface{}) bool, ) // Scan iterates through all data in tree in no specified order. Scan(iter func(min, max [2]float64, data interface{}) bool) // Len returns the number of items in tree Len() int // Bounds returns the minimum bounding box Bounds() (min, max [2]float64) // Children returns all children for parent node. If parent node is nil // then the root nodes should be returned. // The reuse buffer is an empty length slice that can optionally be used // to avoid extra allocations. Children(parent interface{}, reuse []child.Child) (children []child.Child) } // Index is a wrapper around Interface that provides extra features like a // Nearby (kNN) function. // This can be created like such: // var tree = rbang.New() // var index = index.Index{tree} // Now you can use `index` just like tree but with the extra features. type Index struct { tree Interface } // Wrap a tree-like geospatial interface. func Wrap(tree Interface) *Index { return &Index{tree} } // Insert an item into the index func (index *Index) Insert(min, max [2]float64, data interface{}) { index.tree.Insert(min, max, data) } // Search the index for items that intersects the rect param func (index *Index) Search( min, max [2]float64, iter func(min, max [2]float64, data interface{}) bool, ) { index.tree.Search(min, max, iter) } // Delete an item from the index func (index *Index) Delete(min, max [2]float64, data interface{}) { index.tree.Delete(min, max, data) } // Children returns all children for parent node. If parent node is nil // then the root nodes should be returned. // The reuse buffer is an empty length slice that can optionally be used // to avoid extra allocations. func (index *Index) Children(parent interface{}, reuse []child.Child) ( children []child.Child, ) { return index.tree.Children(parent, reuse) } // Nearby performs a kNN-type operation on the index. // It's expected that the caller provides its own the `algo` function, which // is used to calculate a distance to data. The `add` function should be // called by the caller to "return" the data item along with a distance. // 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( 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, ) { var q queue var parent interface{} 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 { // gather all children for parent children = index.tree.Children(parent, children[:0]) for _, child := range children { added = added[:0] algo(child.Min, child.Max, child.Data, child.Item, add) for _, node := range added { q.push(node) } } for { node, ok := q.pop() if !ok { // nothing left in queue return } if node.child.Item { if !iter(node.child.Min, node.child.Max, node.child.Data, node.dist) { return } } else { // gather more children parent = node.child.Data break } } } } // Len returns the number of items in tree func (index *Index) Len() int { return index.tree.Len() } // Bounds returns the minimum bounding box func (index *Index) Bounds() (min, max [2]float64) { return index.tree.Bounds() } // Priority Queue ordered by dist (smallest to largest) type qnode struct { dist float64 child child.Child } type queue struct { nodes []qnode len int size int } func (q *queue) push(node qnode) { if q.nodes == nil { q.nodes = make([]qnode, 2) } else { q.nodes = append(q.nodes, qnode{}) } i := q.len + 1 j := i / 2 for i > 1 && q.nodes[j].dist > node.dist { q.nodes[i] = q.nodes[j] i = j j = j / 2 } q.nodes[i] = node q.len++ } func (q *queue) pop() (qnode, bool) { if q.len == 0 { return qnode{}, false } n := q.nodes[1] q.nodes[1] = q.nodes[q.len] q.len-- var j, k int i := 1 for i != q.len+1 { k = q.len + 1 j = 2 * i if j <= q.len && q.nodes[j].dist < q.nodes[k].dist { k = j } if j+1 <= q.len && q.nodes[j+1].dist < q.nodes[k].dist { k = j + 1 } q.nodes[i] = q.nodes[k] i = k } return n, true } // Scan iterates through all data in tree in no specified order. func (index *Index) Scan( iter func(min, max [2]float64, data interface{}) bool, ) { index.tree.Scan(iter) } // SimpleBoxAlgo performs box-distance algorithm on rectangles. func SimpleBoxAlgo(targetMin, targetMax [2]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, add func(min, max [2]float64, data interface{}, item bool, dist float64), ) { add(min, max, data, item, boxDist(targetMin, targetMax, min, max)) } } func boxDist(amin, amax, bmin, bmax [2]float64) float64 { var dist float64 var min, max float64 if amin[0] > bmin[0] { min = amin[0] } else { min = bmin[0] } if amax[0] < bmax[0] { max = amax[0] } else { max = bmax[0] } squared := min - max if squared > 0 { dist += squared * squared } if amin[1] > bmin[1] { min = amin[1] } else { min = bmin[1] } if amax[1] < bmax[1] { max = amax[1] } else { max = bmax[1] } squared = min - max if squared > 0 { dist += squared * squared } return dist } func (index *Index) svg(child child.Child, height int) []byte { var out []byte point := true for i := 0; i < 2; i++ { if child.Min[i] != child.Max[i] { point = false break } } if point { // is point out = append(out, fmt.Sprintf( "\n", (child.Min[0])*svgScale, (child.Min[1])*svgScale, (child.Max[0]-child.Min[0]+1/svgScale)*svgScale, (child.Max[1]-child.Min[1]+1/svgScale)*svgScale, strokes[height%len(strokes)])...) } else { // is rect out = append(out, fmt.Sprintf( "\n", (child.Min[0])*svgScale, (child.Min[1])*svgScale, (child.Max[0]-child.Min[0]+1/svgScale)*svgScale, (child.Max[1]-child.Min[1]+1/svgScale)*svgScale, strokes[height%len(strokes)])...) } if !child.Item { children := index.tree.Children(child.Data, nil) for _, child := range children { out = append(out, index.svg(child, height+1)...) } } return out } const ( // Continue to first child rectangle and/or next sibling. Continue = iota // Ignore child rectangles but continue to next sibling. Ignore // Stop iterating Stop ) const svgScale = 4.0 var strokes = [...]string{"black", "#cccc00", "green", "red", "purple"} // SVG prints 2D rtree in wgs84 coordinate space func (index *Index) SVG() string { var out string out += fmt.Sprintf("\n", -190.0*svgScale, -100.0*svgScale, 380.0*svgScale, 190.0*svgScale) out += fmt.Sprintf("\n") var outb []byte for _, child := range index.Children(nil, nil) { outb = append(outb, index.svg(child, 1)...) } out += string(outb) out += fmt.Sprintf("\n") out += fmt.Sprintf("\n") return out }