├── .gitignore ├── LICENSE ├── README.md ├── geom.go ├── geom_test.go ├── rtree.go └── rtree_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | *~ 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2012, Daniel Connelly. All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 4 | 5 | Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 6 | 7 | Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. 8 | 9 | Neither the name of Daniel Connelly nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. 10 | 11 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 12 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | rtreego 2 | ======= 3 | 4 | [![GoDoc](https://godoc.org/github.com/patrick-higgins/rtreego?status.svg)](https://godoc.org/github.com/patrick-higgins/rtreego) 5 | 6 | A library for efficiently storing and querying spatial data 7 | in the Go programming language. 8 | 9 | Forked from github.com/dhconnelly/rtreego to specialize for 3 dimensions 10 | and tune for fewer memory allocations. 11 | 12 | 13 | About 14 | ----- 15 | 16 | The R-tree is a popular data structure for efficiently storing and 17 | querying spatial objects; one common use is implementing geospatial 18 | indexes in database management systems. The variant implemented here, 19 | known as the R*-tree, improves performance and increases storage 20 | utilization. Both bounding-box queries and k-nearest-neighbor queries 21 | are supported. 22 | 23 | R-trees are balanced, so maximum tree height is guaranteed to be 24 | logarithmic in the number of entries; however, good worst-case 25 | performance is not guaranteed. Instead, a number of rebalancing 26 | heuristics are applied that perform well in practice. For more 27 | details please refer to the references. 28 | 29 | 30 | Status 31 | ------ 32 | 33 | Geometric primitives (points, rectangles, and their relevant geometric 34 | algorithms) are implemented and tested. The R-tree data structure and 35 | algorithms are currently under development. 36 | 37 | 38 | Install 39 | ------- 40 | 41 | With Go 1 installed, just run `go get github.com/patrick-higgins/rtreego`. 42 | 43 | 44 | Usage 45 | ----- 46 | 47 | Make sure you `import github.com/patrick-higgins/rtreego` in your Go source files. 48 | 49 | ### Storing, updating, and deleting objects 50 | 51 | To create a new tree, specify the number of spatial dimensions and the minimum 52 | and maximum branching factor: 53 | 54 | rt := rtreego.NewTree(2, 25, 50) 55 | 56 | Any type that implements the `Spatial` interface can be stored in the tree: 57 | 58 | type Spatial interface { 59 | Bounds() *Rect 60 | } 61 | 62 | `Rect`s are data structures for representing spatial objects, while `Point`s 63 | represent spatial locations. Creating `Point`s is easy--they're just slices 64 | of `float64`s: 65 | 66 | p1 := rtreego.Point{0.4, 0.5} 67 | p2 := rtreego.Point{6.2, -3.4} 68 | 69 | To create a `Rect`, specify a location and the lengths of the sides: 70 | 71 | r1 := rtreego.NewRect(p1, []float64{1, 2}) 72 | r2 := rtreego.NewRect(p2, []float64{1.7, 2.7}) 73 | 74 | To demonstrate, let's create and store some test data. 75 | 76 | type Thing struct { 77 | where *Rect 78 | name string 79 | } 80 | 81 | func (t *Thing) Bounds() *Rect { 82 | return t.where 83 | } 84 | 85 | rt.Insert(&Thing{r1, "foo"}) 86 | rt.Insert(&Thing{r2, "bar"}) 87 | 88 | size := rt.Size() // returns 2 89 | 90 | We can insert and delete objects from the tree in any order. 91 | 92 | rt.Delete(thing2) 93 | // do some stuff... 94 | rt.Insert(anotherThing) 95 | 96 | If you want to store points instead of rectangles, you can easily convert a 97 | point into a rectangle using the `ToRect` method: 98 | 99 | var tol = 0.01 100 | 101 | type Somewhere struct { 102 | location rtreego.Point 103 | name string 104 | wormhole chan int 105 | } 106 | 107 | func (s *Somewhere) Bounds() *Rect { 108 | // define the bounds of s to be a rectangle centered at s.location 109 | // with side lengths 2 * tol: 110 | return s.location.ToRect(tol) 111 | } 112 | 113 | rt.Insert(&Somewhere{rtreego.Point{0, 0}, "Someplace", nil}) 114 | 115 | If you want to update the location of an object, you must delete it, update it, 116 | and re-insert. Just modifying the object so that the `*Rect` returned by 117 | `Location()` changes, without deleting and re-inserting the object, will 118 | corrupt the tree. 119 | 120 | ### Queries 121 | 122 | Bounding-box and k-nearest-neighbors queries are supported. 123 | 124 | Bounding-box queries require a search `*Rect` argument and come in two flavors: 125 | containment search and intersection search. The former returns all objects that 126 | fall strictly inside the search rectangle, while the latter returns all objects 127 | that touch the search rectangle. 128 | 129 | bb := rtreego.NewRect(rtreego.Point{1.7, -3.4}, []float64{3.2, 1.9}) 130 | 131 | // Get a slice of the objects in rt that intersect bb: 132 | results, _ := rt.SearchIntersect(bb) 133 | 134 | // Get a slice of the objects in rt that are contained inside bb: 135 | results, _ = rt.SearchContained(bb) 136 | 137 | Nearest-neighbor queries find the objects in a tree closest to a specified 138 | query point. 139 | 140 | q := rtreego.Point{6.5, -2.47} 141 | k := 5 142 | 143 | // Get a slice of the k objects in rt closest to q: 144 | results, _ = rt.SearchNearestNeighbors(q, k) 145 | 146 | ### More information 147 | 148 | See http://github.com/patrick-higgins/rtreego for full API documentation. 149 | 150 | 151 | References 152 | ---------- 153 | 154 | - A. Guttman. R-trees: A Dynamic Index Structure for Spatial Searching. 155 | Proceedings of ACM SIGMOD, pages 47-57, 1984. 156 | http://www.cs.jhu.edu/~misha/ReadingSeminar/Papers/Guttman84.pdf 157 | 158 | - N. Beckmann, H .P. Kriegel, R. Schneider and B. Seeger. The R*-tree: An 159 | Efficient and Robust Access Method for Points and Rectangles. Proceedings 160 | of ACM SIGMOD, pages 323-331, May 1990. 161 | http://infolab.usc.edu/csci587/Fall2011/papers/p322-beckmann.pdf 162 | 163 | - N. Roussopoulos, S. Kelley and F. Vincent. Nearest Neighbor Queries. ACM 164 | SIGMOD, pages 71-79, 1995. 165 | http://www.postgis.org/support/nearestneighbor.pdf 166 | 167 | 168 | Author 169 | ------ 170 | 171 | rtreego is written and maintained by Daniel Connelly. You can find my stuff 172 | at dhconnelly.com or email me at dhconnelly@gmail.com. 173 | 174 | This fork is maintained by Patrick Higgins (patrick.allen.higgins@gmail.com). 175 | 176 | 177 | License 178 | ------- 179 | 180 | rtreego is released under a BSD-style license; see LICENSE for more details. 181 | -------------------------------------------------------------------------------- /geom.go: -------------------------------------------------------------------------------- 1 | // Copyright 2012 Daniel Connelly. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package rtreego 6 | 7 | import ( 8 | "fmt" 9 | "math" 10 | "strings" 11 | ) 12 | 13 | // DistError is an improper distance measurement. It implements the error 14 | // and is generated when a distance-related assertion fails. 15 | type DistError float64 16 | 17 | func (err DistError) Error() string { 18 | return "rtreego: improper distance" 19 | } 20 | 21 | // Point represents a point in 3-dimensional Euclidean space. 22 | type Point [Dim]float64 23 | 24 | // Dist computes the Euclidean distance between two points p and q. 25 | func (p Point) dist(q Point) float64 { 26 | sum := 0.0 27 | for i := range p { 28 | dx := p[i] - q[i] 29 | sum += dx * dx 30 | } 31 | return math.Sqrt(sum) 32 | } 33 | 34 | // minDist computes the square of the distance from a point to a rectangle. 35 | // If the point is contained in the rectangle then the distance is zero. 36 | // 37 | // Implemented per Definition 2 of "Nearest Neighbor Queries" by 38 | // N. Roussopoulos, S. Kelley and F. Vincent, ACM SIGMOD, pages 71-79, 1995. 39 | func (p Point) minDist(r *Rect) float64 { 40 | sum := 0.0 41 | for i, pi := range p { 42 | if pi < r.p[i] { 43 | d := pi - r.p[i] 44 | sum += d * d 45 | } else if pi > r.q[i] { 46 | d := pi - r.q[i] 47 | sum += d * d 48 | } else { 49 | sum += 0 50 | } 51 | } 52 | return sum 53 | } 54 | 55 | // minMaxDist computes the minimum of the maximum distances from p to points 56 | // on r. If r is the bounding box of some geometric objects, then there is 57 | // at least one object contained in r within minMaxDist(p, r) of p. 58 | // 59 | // Implemented per Definition 4 of "Nearest Neighbor Queries" by 60 | // N. Roussopoulos, S. Kelley and F. Vincent, ACM SIGMOD, pages 71-79, 1995. 61 | func (p Point) minMaxDist(r *Rect) float64 { 62 | // by definition, MinMaxDist(p, r) = 63 | // min{1<=k<=n}(|pk - rmk|^2 + sum{1<=i<=n, i != k}(|pi - rMi|^2)) 64 | // where rmk and rMk are defined as follows: 65 | 66 | rm := func(k int) float64 { 67 | if p[k] <= (r.p[k]+r.q[k])/2 { 68 | return r.p[k] 69 | } 70 | return r.q[k] 71 | } 72 | 73 | rM := func(k int) float64 { 74 | if p[k] >= (r.p[k]+r.q[k])/2 { 75 | return r.p[k] 76 | } 77 | return r.q[k] 78 | } 79 | 80 | // This formula can be computed in linear time by precomputing 81 | // S = sum{1<=i<=n}(|pi - rMi|^2). 82 | 83 | S := 0.0 84 | for i := range p { 85 | d := p[i] - rM(i) 86 | S += d * d 87 | } 88 | 89 | // Compute MinMaxDist using the precomputed S. 90 | min := math.MaxFloat64 91 | for k := range p { 92 | d1 := p[k] - rM(k) 93 | d2 := p[k] - rm(k) 94 | d := S - d1*d1 + d2*d2 95 | if d < min { 96 | min = d 97 | } 98 | } 99 | 100 | return min 101 | } 102 | 103 | // Rect represents a subset of 3-dimensional Euclidean space of the form 104 | // [a1, b1] x [a2, b2] x ... x [an, bn], where ai < bi for all 1 <= i <= n. 105 | type Rect struct { 106 | p, q Point // Enforced by NewRect: p[i] <= q[i] for all i. 107 | } 108 | 109 | func (r *Rect) String() string { 110 | var s [Dim]string 111 | for i, a := range r.p { 112 | b := r.q[i] 113 | s[i] = fmt.Sprintf("[%.2f, %.2f]", a, b) 114 | } 115 | return strings.Join(s[:], "x") 116 | } 117 | 118 | // NewRect constructs and returns a pointer to a Rect given a corner point and 119 | // the lengths of each dimension. The point p should be the most-negative point 120 | // on the rectangle (in every dimension) and every length should be positive. 121 | func NewRect(p Point, lengths [Dim]float64) (r Rect, err error) { 122 | r.p = p 123 | r.q = lengths 124 | for i, l := range r.q { 125 | if l <= 0 { 126 | return r, DistError(l) 127 | } 128 | r.q[i] += r.p[i] 129 | } 130 | return r, nil 131 | } 132 | 133 | // size computes the measure of a rectangle (the product of its side lengths). 134 | func (r *Rect) size() float64 { 135 | size := 1.0 136 | for i, a := range r.p { 137 | b := r.q[i] 138 | size *= b - a 139 | } 140 | return size 141 | } 142 | 143 | // margin computes the sum of the edge lengths of a rectangle. 144 | func (r *Rect) margin() float64 { 145 | // The number of edges in an n-dimensional rectangle is n * 2^(n-1) 146 | // (http://en.wikipedia.org/wiki/Hypercube_graph). Thus the number 147 | // of edges of length (ai - bi), where the rectangle is determined 148 | // by p = (a1, a2, ..., an) and q = (b1, b2, ..., bn), is 2^(n-1). 149 | // 150 | // The margin of the rectangle, then, is given by the formula 151 | // 2^(n-1) * [(b1 - a1) + (b2 - a2) + ... + (bn - an)]. 152 | sum := 0.0 153 | for i, a := range r.p { 154 | b := r.q[i] 155 | sum += b - a 156 | } 157 | return 4.0 * sum 158 | } 159 | 160 | // containsPoint tests whether p is located inside or on the boundary of r. 161 | func (r *Rect) containsPoint(p Point) bool { 162 | for i, a := range p { 163 | // p is contained in (or on) r if and only if p <= a <= q for 164 | // every dimension. 165 | if a < r.p[i] || a > r.q[i] { 166 | return false 167 | } 168 | } 169 | 170 | return true 171 | } 172 | 173 | // containsRect tests whether r2 is is located inside r1. 174 | func (r1 *Rect) containsRect(r2 *Rect) bool { 175 | for i, a1 := range r1.p { 176 | b1, a2, b2 := r1.q[i], r2.p[i], r2.q[i] 177 | // enforced by constructor: a1 <= b1 and a2 <= b2. 178 | // so containment holds if and only if a1 <= a2 <= b2 <= b1 179 | // for every dimension. 180 | if a1 > a2 || b2 > b1 { 181 | return false 182 | } 183 | } 184 | 185 | return true 186 | } 187 | 188 | func (r1 *Rect) enlarge(r2 *Rect) { 189 | for i := 0; i < Dim; i++ { 190 | if r1.p[i] > r2.p[i] { 191 | r1.p[i] = r2.p[i] 192 | } 193 | if r1.q[i] < r2.q[i] { 194 | r1.q[i] = r2.q[i] 195 | } 196 | } 197 | } 198 | 199 | // intersect computes the intersection of two rectangles. If no intersection 200 | // exists, the intersection is nil. 201 | func intersect(r1, r2 *Rect) bool { 202 | // There are four cases of overlap: 203 | // 204 | // 1. a1------------b1 205 | // a2------------b2 206 | // p--------q 207 | // 208 | // 2. a1------------b1 209 | // a2------------b2 210 | // p--------q 211 | // 212 | // 3. a1-----------------b1 213 | // a2-------b2 214 | // p--------q 215 | // 216 | // 4. a1-------b1 217 | // a2-----------------b2 218 | // p--------q 219 | // 220 | // Thus there are only two cases of non-overlap: 221 | // 222 | // 1. a1------b1 223 | // a2------b2 224 | // 225 | // 2. a1------b1 226 | // a2------b2 227 | // 228 | // Enforced by constructor: a1 <= b1 and a2 <= b2. So we can just 229 | // check the endpoints. 230 | 231 | for i := 0; i < Dim; i++ { 232 | if r2.q[i] <= r1.p[i] || r1.q[i] <= r2.p[i] { 233 | return false 234 | } 235 | } 236 | return true 237 | } 238 | 239 | // ToRect constructs a rectangle containing p with side lengths 2*tol. 240 | func (p Point) ToRect(tol float64) *Rect { 241 | var r Rect 242 | for i := range p { 243 | r.p[i] = p[i] - tol 244 | r.q[i] = p[i] + tol 245 | } 246 | return &r 247 | } 248 | 249 | func initBoundingBox(r, r1, r2 *Rect) { 250 | *r = *r1 251 | r.enlarge(r2) 252 | } 253 | 254 | // boundingBox constructs the smallest rectangle containing both r1 and r2. 255 | func boundingBox(r1, r2 *Rect) *Rect { 256 | var r Rect 257 | initBoundingBox(&r, r1, r2) 258 | return &r 259 | } 260 | -------------------------------------------------------------------------------- /geom_test.go: -------------------------------------------------------------------------------- 1 | // Copyright 2012 Daniel Connelly. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | package rtreego 6 | 7 | import ( 8 | "math" 9 | "testing" 10 | ) 11 | 12 | const EPS = 0.000000001 13 | 14 | func TestDist(t *testing.T) { 15 | p := Point{1, 2, 3} 16 | q := Point{4, 5, 6} 17 | dist := math.Sqrt(27) 18 | if d := p.dist(q); d != dist { 19 | t.Errorf("dist(%v, %v) = %v; expected %v", p, q, d, dist) 20 | } 21 | } 22 | 23 | func TestNewRect(t *testing.T) { 24 | p := Point{1.0, -2.5, 3.0} 25 | q := Point{3.5, 5.5, 4.5} 26 | lengths := [Dim]float64{2.5, 8.0, 1.5} 27 | 28 | rect, err := NewRect(p, lengths) 29 | if err != nil { 30 | t.Errorf("Error on NewRect(%v, %v): %v", p, lengths, err) 31 | } 32 | if d := p.dist(rect.p); d > EPS { 33 | t.Errorf("Expected p == rect.p") 34 | } 35 | if d := q.dist(rect.q); d > EPS { 36 | t.Errorf("Expected q == rect.q") 37 | } 38 | } 39 | 40 | func TestNewRectDistError(t *testing.T) { 41 | p := Point{1.0, -2.5, 3.0} 42 | lengths := [Dim]float64{2.5, -8.0, 1.5} 43 | _, err := NewRect(p, lengths) 44 | if _, ok := err.(DistError); !ok { 45 | t.Errorf("Expected distError on NewRect(%v, %v)", p, lengths) 46 | } 47 | } 48 | 49 | func TestRectSize(t *testing.T) { 50 | p := Point{1.0, -2.5, 3.0} 51 | lengths := [Dim]float64{2.5, 8.0, 1.5} 52 | rect, _ := NewRect(p, lengths) 53 | size := lengths[0] * lengths[1] * lengths[2] 54 | actual := rect.size() 55 | if size != actual { 56 | t.Errorf("Expected %v.size() == %v, got %v", rect, size, actual) 57 | } 58 | } 59 | 60 | func TestRectMargin(t *testing.T) { 61 | p := Point{1.0, -2.5, 3.0} 62 | lengths := [Dim]float64{2.5, 8.0, 1.5} 63 | rect, _ := NewRect(p, lengths) 64 | size := 4*2.5 + 4*8.0 + 4*1.5 65 | actual := rect.margin() 66 | if size != actual { 67 | t.Errorf("Expected %v.margin() == %v, got %v", rect, size, actual) 68 | } 69 | } 70 | 71 | func TestContainsPoint(t *testing.T) { 72 | p := Point{3.7, -2.4, 0.0} 73 | lengths := [Dim]float64{6.2, 1.1, 4.9} 74 | rect, _ := NewRect(p, lengths) 75 | 76 | q := Point{4.5, -1.7, 4.8} 77 | if yes := rect.containsPoint(q); !yes { 78 | t.Errorf("Expected %v contains %v", rect, q) 79 | } 80 | } 81 | 82 | func TestDoesNotContainPoint(t *testing.T) { 83 | p := Point{3.7, -2.4, 0.0} 84 | lengths := [Dim]float64{6.2, 1.1, 4.9} 85 | rect, _ := NewRect(p, lengths) 86 | 87 | q := Point{4.5, -1.7, -3.2} 88 | if yes := rect.containsPoint(q); yes { 89 | t.Errorf("Expected %v doesn't contain %v", rect, q) 90 | } 91 | } 92 | 93 | func TestContainsRect(t *testing.T) { 94 | p := Point{3.7, -2.4, 0.0} 95 | lengths1 := [Dim]float64{6.2, 1.1, 4.9} 96 | rect1, _ := NewRect(p, lengths1) 97 | 98 | q := Point{4.1, -1.9, 1.0} 99 | lengths2 := [Dim]float64{3.2, 0.6, 3.7} 100 | rect2, _ := NewRect(q, lengths2) 101 | 102 | if yes := rect1.containsRect(&rect2); !yes { 103 | t.Errorf("Expected %v.containsRect(%v", rect1, rect2) 104 | } 105 | } 106 | 107 | func TestDoesNotContainRectOverlaps(t *testing.T) { 108 | p := Point{3.7, -2.4, 0.0} 109 | lengths1 := [Dim]float64{6.2, 1.1, 4.9} 110 | rect1, _ := NewRect(p, lengths1) 111 | 112 | q := Point{4.1, -1.9, 1.0} 113 | lengths2 := [Dim]float64{3.2, 1.4, 3.7} 114 | rect2, _ := NewRect(q, lengths2) 115 | 116 | if yes := rect1.containsRect(&rect2); yes { 117 | t.Errorf("Expected %v doesn't contain %v", rect1, rect2) 118 | } 119 | } 120 | 121 | func TestDoesNotContainRectDisjoint(t *testing.T) { 122 | p := Point{3.7, -2.4, 0.0} 123 | lengths1 := [Dim]float64{6.2, 1.1, 4.9} 124 | rect1, _ := NewRect(p, lengths1) 125 | 126 | q := Point{1.2, -19.6, -4.0} 127 | lengths2 := [Dim]float64{2.2, 5.9, 0.5} 128 | rect2, _ := NewRect(q, lengths2) 129 | 130 | if yes := rect1.containsRect(&rect2); yes { 131 | t.Errorf("Expected %v doesn't contain %v", rect1, rect2) 132 | } 133 | } 134 | 135 | func TestNoIntersection(t *testing.T) { 136 | p := Point{1, 2, 3} 137 | lengths1 := [Dim]float64{1, 1, 1} 138 | rect1, _ := NewRect(p, lengths1) 139 | 140 | q := Point{-1, -2, -3} 141 | lengths2 := [Dim]float64{2.5, 3, 6.5} 142 | rect2, _ := NewRect(q, lengths2) 143 | 144 | // rect1 and rect2 fail to overlap in just one dimension (second) 145 | 146 | if intersect(&rect1, &rect2) { 147 | t.Errorf("Expected intersect(%v, %v) == false", rect1, rect2) 148 | } 149 | } 150 | 151 | func TestNoIntersectionJustTouches(t *testing.T) { 152 | p := Point{1, 2, 3} 153 | lengths1 := [Dim]float64{1, 1, 1} 154 | rect1, _ := NewRect(p, lengths1) 155 | 156 | q := Point{-1, -2, -3} 157 | lengths2 := [Dim]float64{2.5, 4, 6.5} 158 | rect2, _ := NewRect(q, lengths2) 159 | 160 | // rect1 and rect2 fail to overlap in just one dimension (second) 161 | 162 | if intersect(&rect1, &rect2) { 163 | t.Errorf("Expected intersect(%v, %v) == nil", rect1, rect2) 164 | } 165 | } 166 | 167 | func TestContainmentIntersection(t *testing.T) { 168 | p := Point{1, 2, 3} 169 | lengths1 := [Dim]float64{1, 1, 1} 170 | rect1, _ := NewRect(p, lengths1) 171 | 172 | q := Point{1, 2.2, 3.3} 173 | lengths2 := [Dim]float64{0.5, 0.5, 0.5} 174 | rect2, _ := NewRect(q, lengths2) 175 | 176 | r := Point{1, 2.2, 3.3} 177 | s := Point{1.5, 2.7, 3.8} 178 | 179 | if !intersect(&rect1, &rect2) { 180 | t.Errorf("intersect(%v, %v) != %v, %v", rect1, rect2, r, s) 181 | } 182 | } 183 | 184 | func TestOverlapIntersection(t *testing.T) { 185 | p := Point{1, 2, 3} 186 | lengths1 := [Dim]float64{1, 2.5, 1} 187 | rect1, _ := NewRect(p, lengths1) 188 | 189 | q := Point{1, 4, -3} 190 | lengths2 := [Dim]float64{3, 2, 6.5} 191 | rect2, _ := NewRect(q, lengths2) 192 | 193 | r := Point{1, 4, 3} 194 | s := Point{2, 4.5, 3.5} 195 | 196 | if !intersect(&rect1, &rect2) { 197 | t.Errorf("intersect(%v, %v) != %v, %v", rect1, rect2, r, s) 198 | } 199 | } 200 | 201 | func TestToRect(t *testing.T) { 202 | x := Point{3.7, -2.4, 0.0} 203 | tol := 0.05 204 | rect := x.ToRect(tol) 205 | 206 | p := Point{3.65, -2.45, -0.05} 207 | q := Point{3.75, -2.35, 0.05} 208 | d1 := p.dist(rect.p) 209 | d2 := q.dist(rect.q) 210 | if d1 > EPS || d2 > EPS { 211 | t.Errorf("Expected %v.ToRect(%v) == %v, %v, got %v", x, tol, p, q, rect) 212 | } 213 | } 214 | 215 | func TestBoundingBox(t *testing.T) { 216 | p := Point{3.7, -2.4, 0.0} 217 | lengths1 := [Dim]float64{1, 15, 3} 218 | rect1, _ := NewRect(p, lengths1) 219 | 220 | q := Point{-6.5, 4.7, 2.5} 221 | lengths2 := [Dim]float64{4, 5, 6} 222 | rect2, _ := NewRect(q, lengths2) 223 | 224 | r := Point{-6.5, -2.4, 0.0} 225 | s := Point{4.7, 12.6, 8.5} 226 | 227 | bb := boundingBox(&rect1, &rect2) 228 | d1 := r.dist(bb.p) 229 | d2 := s.dist(bb.q) 230 | if d1 > EPS || d2 > EPS { 231 | t.Errorf("boundingBox(%v, %v) != %v, %v, got %v", rect1, rect2, r, s, bb) 232 | } 233 | } 234 | 235 | func TestBoundingBoxContains(t *testing.T) { 236 | p := Point{3.7, -2.4, 0.0} 237 | lengths1 := [Dim]float64{1, 15, 3} 238 | rect1, _ := NewRect(p, lengths1) 239 | 240 | q := Point{4.0, 0.0, 1.5} 241 | lengths2 := [Dim]float64{0.56, 6.222222, 0.946} 242 | rect2, _ := NewRect(q, lengths2) 243 | 244 | bb := boundingBox(&rect1, &rect2) 245 | d1 := rect1.p.dist(bb.p) 246 | d2 := rect1.q.dist(bb.q) 247 | if d1 > EPS || d2 > EPS { 248 | t.Errorf("boundingBox(%v, %v) != %v, got %v", rect1, rect2, rect1, bb) 249 | } 250 | } 251 | 252 | func TestMinDistZero(t *testing.T) { 253 | p := Point{1, 2, 3} 254 | r := p.ToRect(1) 255 | if d := p.minDist(r); d > EPS { 256 | t.Errorf("Expected %v.minDist(%v) == 0, got %v", p, r, d) 257 | } 258 | } 259 | 260 | func TestMinDistPositive(t *testing.T) { 261 | p := Point{1, 2, 3} 262 | r := Rect{Point{-1, -4, 7}, Point{2, -2, 9}} 263 | expected := float64((-2-2)*(-2-2) + (7-3)*(7-3)) 264 | if d := p.minDist(&r); math.Abs(d-expected) > EPS { 265 | t.Errorf("Expected %v.minDist(%v) == %v, got %v", p, r, expected, d) 266 | } 267 | } 268 | 269 | func TestMinMaxdist(t *testing.T) { 270 | p := Point{-3, -2, -1} 271 | r := Rect{Point{0, 0, 0}, Point{1, 2, 3}} 272 | 273 | // furthest points from p on the faces closest to p in each dimension 274 | q1 := Point{0, 2, 3} 275 | q2 := Point{1, 0, 3} 276 | q3 := Point{1, 2, 0} 277 | 278 | // find the closest distance from p to one of these furthest points 279 | d1 := p.dist(q1) 280 | d2 := p.dist(q2) 281 | d3 := p.dist(q3) 282 | expected := math.Min(d1*d1, math.Min(d2*d2, d3*d3)) 283 | 284 | if d := p.minMaxDist(&r); math.Abs(d-expected) > EPS { 285 | t.Errorf("Expected %v.minMaxDist(%v) == %v, got %v", p, r, expected, d) 286 | } 287 | } 288 | -------------------------------------------------------------------------------- /rtree.go: -------------------------------------------------------------------------------- 1 | // Copyright 2012 Daniel Connelly. All rights reserved. 2 | // Use of this source code is governed by a BSD-style 3 | // license that can be found in the LICENSE file. 4 | 5 | // A library for efficiently storing and querying spatial data. 6 | package rtreego 7 | 8 | import ( 9 | "fmt" 10 | "math" 11 | "sort" 12 | ) 13 | 14 | const Dim = 3 15 | 16 | // Rtree represents an R-tree, a balanced search tree for storing and querying 17 | // spatial objects. MinChildren/MaxChildren specify the minimum/maximum branching factors. 18 | type Rtree struct { 19 | MinChildren int 20 | MaxChildren int 21 | root *node 22 | size int 23 | height int 24 | } 25 | 26 | // NewTree creates a new R-tree instance. 27 | func NewTree(MinChildren, MaxChildren int) *Rtree { 28 | rt := Rtree{MinChildren: MinChildren, MaxChildren: MaxChildren} 29 | rt.height = 1 30 | rt.root = &node{} 31 | rt.root.entries = make([]entry, 0, MaxChildren) 32 | rt.root.leaf = true 33 | rt.root.level = 1 34 | return &rt 35 | } 36 | 37 | // Size returns the number of objects currently stored in tree. 38 | func (tree *Rtree) Size() int { 39 | return tree.size 40 | } 41 | 42 | func (tree *Rtree) String() string { 43 | return "(*Rtree)" 44 | } 45 | 46 | // Depth returns the maximum depth of tree. 47 | func (tree *Rtree) Depth() int { 48 | return tree.height 49 | } 50 | 51 | // node represents a tree node of an Rtree. 52 | type node struct { 53 | parent *node 54 | leaf bool 55 | entries []entry 56 | level int // node depth in the Rtree 57 | } 58 | 59 | func (n *node) String() string { 60 | return fmt.Sprintf("node{leaf: %v, entries: %v}", n.leaf, n.entries) 61 | } 62 | 63 | // entry represents a spatial index record stored in a tree node. 64 | type entry struct { 65 | bb *Rect // bounding-box of all children of this entry 66 | child *node 67 | obj Spatial 68 | } 69 | 70 | func (e entry) String() string { 71 | if e.child != nil { 72 | return fmt.Sprintf("entry{bb: %v, child: %v}", e.bb, e.child) 73 | } 74 | return fmt.Sprintf("entry{bb: %v, obj: %v}", e.bb, e.obj) 75 | } 76 | 77 | // Any type that implements Spatial can be stored in an Rtree and queried. 78 | type Spatial interface { 79 | Bounds() *Rect 80 | } 81 | 82 | // Insertion 83 | 84 | // Insert inserts a spatial object into the tree. If insertion 85 | // causes a leaf node to overflow, the tree is rebalanced automatically. 86 | // 87 | // Implemented per Section 3.2 of "R-trees: A Dynamic Index Structure for 88 | // Spatial Searching" by A. Guttman, Proceedings of ACM SIGMOD, p. 47-57, 1984. 89 | func (tree *Rtree) Insert(obj Spatial) { 90 | e := entry{obj.Bounds(), nil, obj} 91 | tree.insert(e, 1) 92 | tree.size++ 93 | } 94 | 95 | // insert adds the specified entry to the tree at the specified level. 96 | func (tree *Rtree) insert(e entry, level int) { 97 | leaf := tree.chooseNode(tree.root, e, level) 98 | leaf.entries = append(leaf.entries, e) 99 | 100 | // update parent pointer if necessary 101 | if e.child != nil { 102 | e.child.parent = leaf 103 | } 104 | 105 | // split leaf if overflows 106 | var split *node 107 | if len(leaf.entries) > tree.MaxChildren { 108 | leaf, split = leaf.split(tree.MinChildren) 109 | } 110 | root, splitRoot := tree.adjustTree(leaf, split) 111 | if splitRoot != nil { 112 | oldRoot := root 113 | tree.height++ 114 | tree.root = &node{ 115 | parent: nil, 116 | level: tree.height, 117 | entries: []entry{ 118 | entry{bb: oldRoot.computeBoundingBox(), child: oldRoot}, 119 | entry{bb: splitRoot.computeBoundingBox(), child: splitRoot}, 120 | }, 121 | } 122 | oldRoot.parent = tree.root 123 | splitRoot.parent = tree.root 124 | } 125 | } 126 | 127 | // chooseNode finds the node at the specified level to which e should be added. 128 | func (tree *Rtree) chooseNode(n *node, e entry, level int) *node { 129 | if n.leaf || n.level == level { 130 | return n 131 | } 132 | 133 | // find the entry whose bb needs least enlargement to include obj 134 | diff := math.MaxFloat64 135 | var chosen entry 136 | var bb Rect 137 | for _, en := range n.entries { 138 | initBoundingBox(&bb, en.bb, e.bb) 139 | d := bb.size() - en.bb.size() 140 | if d < diff || (d == diff && en.bb.size() < chosen.bb.size()) { 141 | diff = d 142 | chosen = en 143 | } 144 | } 145 | 146 | return tree.chooseNode(chosen.child, e, level) 147 | } 148 | 149 | // adjustTree splits overflowing nodes and propagates the changes upwards. 150 | func (tree *Rtree) adjustTree(n, nn *node) (*node, *node) { 151 | // Let the caller handle root adjustments. 152 | if n == tree.root { 153 | return n, nn 154 | } 155 | 156 | // Re-size the bounding box of n to account for lower-level changes. 157 | en := n.getEntry() 158 | en.bb = n.computeBoundingBox() 159 | 160 | // If nn is nil, then we're just propagating changes upwards. 161 | if nn == nil { 162 | return tree.adjustTree(n.parent, nil) 163 | } 164 | 165 | // Otherwise, these are two nodes resulting from a split. 166 | // n was reused as the "left" node, but we need to add nn to n.parent. 167 | enn := entry{nn.computeBoundingBox(), nn, nil} 168 | n.parent.entries = append(n.parent.entries, enn) 169 | 170 | // If the new entry overflows the parent, split the parent and propagate. 171 | if len(n.parent.entries) > tree.MaxChildren { 172 | return tree.adjustTree(n.parent.split(tree.MinChildren)) 173 | } 174 | 175 | // Otherwise keep propagating changes upwards. 176 | return tree.adjustTree(n.parent, nil) 177 | } 178 | 179 | // getEntry returns a pointer to the entry for the node n from n's parent. 180 | func (n *node) getEntry() *entry { 181 | var e *entry 182 | for i := range n.parent.entries { 183 | if n.parent.entries[i].child == n { 184 | e = &n.parent.entries[i] 185 | break 186 | } 187 | } 188 | return e 189 | } 190 | 191 | // computeBoundingBox finds the MBR of the children of n. 192 | func (n *node) computeBoundingBox() *Rect { 193 | var bb Rect 194 | for i, e := range n.entries { 195 | if i == 0 { 196 | bb = *e.bb 197 | } else { 198 | bb.enlarge(e.bb) 199 | } 200 | } 201 | return &bb 202 | } 203 | 204 | // split splits a node into two groups while attempting to minimize the 205 | // bounding-box area of the resulting groups. 206 | func (n *node) split(minGroupSize int) (left, right *node) { 207 | // find the initial split 208 | l, r := n.pickSeeds() 209 | leftSeed, rightSeed := n.entries[l], n.entries[r] 210 | 211 | // get the entries to be divided between left and right 212 | remaining := append(n.entries[:l], n.entries[l+1:r]...) 213 | remaining = append(remaining, n.entries[r+1:]...) 214 | 215 | // setup the new split nodes, but re-use n as the left node 216 | left = n 217 | left.entries = []entry{leftSeed} 218 | right = &node{ 219 | parent: n.parent, 220 | leaf: n.leaf, 221 | level: n.level, 222 | entries: []entry{rightSeed}, 223 | } 224 | 225 | // TODO 226 | if rightSeed.child != nil { 227 | rightSeed.child.parent = right 228 | } 229 | if leftSeed.child != nil { 230 | leftSeed.child.parent = left 231 | } 232 | 233 | // distribute all of n's old entries into left and right. 234 | for len(remaining) > 0 { 235 | next := pickNext(left, right, remaining) 236 | e := remaining[next] 237 | 238 | if len(remaining)+len(left.entries) <= minGroupSize { 239 | assign(e, left) 240 | } else if len(remaining)+len(right.entries) <= minGroupSize { 241 | assign(e, right) 242 | } else { 243 | assignGroup(e, left, right) 244 | } 245 | 246 | remaining = append(remaining[:next], remaining[next+1:]...) 247 | } 248 | 249 | return 250 | } 251 | 252 | func assign(e entry, group *node) { 253 | if e.child != nil { 254 | e.child.parent = group 255 | } 256 | group.entries = append(group.entries, e) 257 | } 258 | 259 | // assignGroup chooses one of two groups to which a node should be added. 260 | func assignGroup(e entry, left, right *node) { 261 | leftBB := left.computeBoundingBox() 262 | rightBB := right.computeBoundingBox() 263 | leftEnlarged := boundingBox(leftBB, e.bb) 264 | rightEnlarged := boundingBox(rightBB, e.bb) 265 | 266 | // first, choose the group that needs the least enlargement 267 | leftDiff := leftEnlarged.size() - leftBB.size() 268 | rightDiff := rightEnlarged.size() - rightBB.size() 269 | if diff := leftDiff - rightDiff; diff < 0 { 270 | assign(e, left) 271 | return 272 | } else if diff > 0 { 273 | assign(e, right) 274 | return 275 | } 276 | 277 | // next, choose the group that has smaller area 278 | if diff := leftBB.size() - rightBB.size(); diff < 0 { 279 | assign(e, left) 280 | return 281 | } else if diff > 0 { 282 | assign(e, right) 283 | return 284 | } 285 | 286 | // next, choose the group with fewer entries 287 | if diff := len(left.entries) - len(right.entries); diff <= 0 { 288 | assign(e, left) 289 | return 290 | } 291 | assign(e, right) 292 | } 293 | 294 | // pickSeeds chooses two child entries of n to start a split. 295 | func (n *node) pickSeeds() (int, int) { 296 | left, right := 0, 1 297 | maxWastedSpace := -1.0 298 | var bb Rect 299 | for i, e1 := range n.entries { 300 | for j, e2 := range n.entries[i+1:] { 301 | initBoundingBox(&bb, e1.bb, e2.bb) 302 | d := bb.size() - e1.bb.size() - e2.bb.size() 303 | if d > maxWastedSpace { 304 | maxWastedSpace = d 305 | left, right = i, j+i+1 306 | } 307 | } 308 | } 309 | return left, right 310 | } 311 | 312 | // pickNext chooses an entry to be added to an entry group. 313 | func pickNext(left, right *node, entries []entry) (next int) { 314 | maxDiff := -1.0 315 | leftBB := left.computeBoundingBox() 316 | rightBB := right.computeBoundingBox() 317 | for i, e := range entries { 318 | d1 := boundingBox(leftBB, e.bb).size() - leftBB.size() 319 | d2 := boundingBox(rightBB, e.bb).size() - rightBB.size() 320 | d := math.Abs(d1 - d2) 321 | if d > maxDiff { 322 | maxDiff = d 323 | next = i 324 | } 325 | } 326 | return 327 | } 328 | 329 | // Deletion 330 | 331 | // Delete removes an object from the tree. If the object is not found, ok 332 | // is false; otherwise ok is true. 333 | // 334 | // Implemented per Section 3.3 of "R-trees: A Dynamic Index Structure for 335 | // Spatial Searching" by A. Guttman, Proceedings of ACM SIGMOD, p. 47-57, 1984. 336 | func (tree *Rtree) Delete(obj Spatial) bool { 337 | n := tree.findLeaf(tree.root, obj) 338 | if n == nil { 339 | return false 340 | } 341 | 342 | ind := -1 343 | for i, e := range n.entries { 344 | if e.obj == obj { 345 | ind = i 346 | } 347 | } 348 | if ind < 0 { 349 | return false 350 | } 351 | 352 | n.entries = append(n.entries[:ind], n.entries[ind+1:]...) 353 | 354 | tree.condenseTree(n) 355 | tree.size-- 356 | 357 | if !tree.root.leaf && len(tree.root.entries) == 1 { 358 | tree.root = tree.root.entries[0].child 359 | } 360 | 361 | return true 362 | } 363 | 364 | // findLeaf finds the leaf node containing obj. 365 | func (tree *Rtree) findLeaf(n *node, obj Spatial) *node { 366 | if n.leaf { 367 | return n 368 | } 369 | // if not leaf, search all candidate subtrees 370 | for _, e := range n.entries { 371 | if e.bb.containsRect(obj.Bounds()) { 372 | leaf := tree.findLeaf(e.child, obj) 373 | if leaf == nil { 374 | continue 375 | } 376 | // check if the leaf actually contains the object 377 | for _, leafEntry := range leaf.entries { 378 | if leafEntry.obj == obj { 379 | return leaf 380 | } 381 | } 382 | } 383 | } 384 | return nil 385 | } 386 | 387 | // condenseTree deletes underflowing nodes and propagates the changes upwards. 388 | func (tree *Rtree) condenseTree(n *node) { 389 | deleted := []*node{} 390 | 391 | for n != tree.root { 392 | if len(n.entries) < tree.MinChildren { 393 | // remove n from parent entries 394 | entries := []entry{} 395 | for _, e := range n.parent.entries { 396 | if e.child != n { 397 | entries = append(entries, e) 398 | } 399 | } 400 | if len(n.parent.entries) == len(entries) { 401 | panic(fmt.Errorf("Failed to remove entry from parent")) 402 | } 403 | n.parent.entries = entries 404 | 405 | // only add n to deleted if it still has children 406 | if len(n.entries) > 0 { 407 | deleted = append(deleted, n) 408 | } 409 | } else { 410 | // just a child entry deletion, no underflow 411 | n.getEntry().bb = n.computeBoundingBox() 412 | } 413 | n = n.parent 414 | } 415 | 416 | for _, n := range deleted { 417 | // reinsert entry so that it will remain at the same level as before 418 | e := entry{n.computeBoundingBox(), n, nil} 419 | tree.insert(e, n.level+1) 420 | } 421 | } 422 | 423 | // Searching 424 | 425 | // SearchIntersectBB returns all objects that intersect the specified rectangle. 426 | // 427 | // Implemented per Section 3.1 of "R-trees: A Dynamic Index Structure for 428 | // Spatial Searching" by A. Guttman, Proceedings of ACM SIGMOD, p. 47-57, 1984. 429 | func (tree *Rtree) SearchIntersect(bb *Rect) []Spatial { 430 | results := []Spatial{} 431 | return tree.searchIntersect(tree.root, bb, results) 432 | } 433 | 434 | func (tree *Rtree) searchIntersect(n *node, bb *Rect, results []Spatial) []Spatial { 435 | for _, e := range n.entries { 436 | if intersect(e.bb, bb) { 437 | if n.leaf { 438 | results = append(results, e.obj) 439 | } else { 440 | results = tree.searchIntersect(e.child, bb, results) 441 | } 442 | } 443 | } 444 | return results 445 | } 446 | 447 | // NearestNeighbor returns the closest object to the specified point. 448 | // Implemented per "Nearest Neighbor Queries" by Roussopoulos et al 449 | func (tree *Rtree) NearestNeighbor(p Point) Spatial { 450 | obj, _ := tree.nearestNeighbor(p, tree.root, math.MaxFloat64, nil) 451 | return obj 452 | } 453 | 454 | // utilities for sorting slices of entries 455 | 456 | type entrySlice struct { 457 | entries []entry 458 | dists []float64 459 | pt Point 460 | } 461 | 462 | func (s entrySlice) Len() int { return len(s.entries) } 463 | 464 | func (s entrySlice) Swap(i, j int) { 465 | s.entries[i], s.entries[j] = s.entries[j], s.entries[i] 466 | s.dists[i], s.dists[j] = s.dists[j], s.dists[i] 467 | } 468 | 469 | func (s entrySlice) Less(i, j int) bool { 470 | return s.dists[i] < s.dists[j] 471 | } 472 | 473 | func sortEntries(p Point, entries []entry) ([]entry, []float64) { 474 | sorted := make([]entry, len(entries)) 475 | dists := make([]float64, len(entries)) 476 | for i := 0; i < len(entries); i++ { 477 | sorted[i] = entries[i] 478 | dists[i] = p.minDist(entries[i].bb) 479 | } 480 | sort.Sort(entrySlice{sorted, dists, p}) 481 | return sorted, dists 482 | } 483 | 484 | func pruneEntries(p Point, entries []entry, minDists []float64) []entry { 485 | minMinMaxDist := math.MaxFloat64 486 | for i := range entries { 487 | minMaxDist := p.minMaxDist(entries[i].bb) 488 | if minMaxDist < minMinMaxDist { 489 | minMinMaxDist = minMaxDist 490 | } 491 | } 492 | // remove all entries with minDist > minMinMaxDist 493 | pruned := []entry{} 494 | for i := range entries { 495 | if minDists[i] <= minMinMaxDist { 496 | pruned = append(pruned, entries[i]) 497 | } 498 | } 499 | return pruned 500 | } 501 | 502 | func (tree *Rtree) nearestNeighbor(p Point, n *node, d float64, nearest Spatial) (Spatial, float64) { 503 | if n.leaf { 504 | for _, e := range n.entries { 505 | dist := math.Sqrt(p.minDist(e.bb)) 506 | if dist < d { 507 | d = dist 508 | nearest = e.obj 509 | } 510 | } 511 | } else { 512 | branches, dists := sortEntries(p, n.entries) 513 | branches = pruneEntries(p, branches, dists) 514 | for _, e := range branches { 515 | subNearest, dist := tree.nearestNeighbor(p, e.child, d, nearest) 516 | if dist < d { 517 | d = dist 518 | nearest = subNearest 519 | } 520 | } 521 | } 522 | 523 | return nearest, d 524 | } 525 | 526 | func (tree *Rtree) NearestNeighbors(k int, p Point) []Spatial { 527 | dists := make([]float64, k) 528 | objs := make([]Spatial, k) 529 | for i := 0; i < k; i++ { 530 | dists[i] = math.MaxFloat64 531 | objs[i] = nil 532 | } 533 | objs, _ = tree.nearestNeighbors(k, p, tree.root, dists, objs) 534 | return objs 535 | } 536 | 537 | // insert obj into nearest and return the first k elements in increasing order. 538 | func insertNearest(k int, dists []float64, nearest []Spatial, dist float64, obj Spatial) ([]float64, []Spatial) { 539 | i := 0 540 | for i < k && dist >= dists[i] { 541 | i++ 542 | } 543 | if i >= k { 544 | return dists, nearest 545 | } 546 | 547 | left, right := dists[:i], dists[i:k-1] 548 | updatedDists := make([]float64, k) 549 | copy(updatedDists, left) 550 | updatedDists[i] = dist 551 | copy(updatedDists[i+1:], right) 552 | 553 | leftObjs, rightObjs := nearest[:i], nearest[i:k-1] 554 | updatedNearest := make([]Spatial, k) 555 | copy(updatedNearest, leftObjs) 556 | updatedNearest[i] = obj 557 | copy(updatedNearest[i+1:], rightObjs) 558 | 559 | return updatedDists, updatedNearest 560 | } 561 | 562 | func (tree *Rtree) nearestNeighbors(k int, p Point, n *node, dists []float64, nearest []Spatial) ([]Spatial, []float64) { 563 | if n.leaf { 564 | for _, e := range n.entries { 565 | dist := math.Sqrt(p.minDist(e.bb)) 566 | dists, nearest = insertNearest(k, dists, nearest, dist, e.obj) 567 | } 568 | } else { 569 | branches, branchDists := sortEntries(p, n.entries) 570 | branches = pruneEntries(p, branches, branchDists) 571 | for _, e := range branches { 572 | nearest, dists = tree.nearestNeighbors(k, p, e.child, dists, nearest) 573 | } 574 | } 575 | return nearest, dists 576 | } 577 | -------------------------------------------------------------------------------- /rtree_test.go: -------------------------------------------------------------------------------- 1 | package rtreego 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "strings" 7 | "testing" 8 | ) 9 | 10 | func (r *Rect) Bounds() *Rect { 11 | return r 12 | } 13 | 14 | func mustRect(p Point, widths [Dim]float64) *Rect { 15 | if widths[Dim-1] == 0 { 16 | widths[Dim-1] = 1 17 | } 18 | r, err := NewRect(p, widths) 19 | if err != nil { 20 | panic(err) 21 | } 22 | return &r 23 | } 24 | 25 | func printNode(n *node, level int) { 26 | padding := strings.Repeat("\t", level) 27 | fmt.Printf("%sNode: %p\n", padding, n) 28 | fmt.Printf("%sParent: %p\n", padding, n.parent) 29 | fmt.Printf("%sLevel: %d\n", padding, n.level) 30 | fmt.Printf("%sLeaf: %t\n%sEntries:\n", padding, n.leaf, padding) 31 | for _, e := range n.entries { 32 | printEntry(e, level+1) 33 | } 34 | } 35 | 36 | func printEntry(e entry, level int) { 37 | padding := strings.Repeat("\t", level) 38 | fmt.Printf("%sBB: %v\n", padding, e.bb) 39 | if e.child != nil { 40 | printNode(e.child, level) 41 | } else { 42 | fmt.Printf("%sObject: %p\n", padding, e.obj) 43 | } 44 | fmt.Println() 45 | } 46 | 47 | func items(n *node) chan Spatial { 48 | ch := make(chan Spatial) 49 | go func() { 50 | for _, e := range n.entries { 51 | if n.leaf { 52 | ch <- e.obj 53 | } else { 54 | for obj := range items(e.child) { 55 | ch <- obj 56 | } 57 | } 58 | } 59 | close(ch) 60 | }() 61 | return ch 62 | } 63 | 64 | func verify(t *testing.T, n *node) { 65 | if n.leaf { 66 | return 67 | } 68 | for _, e := range n.entries { 69 | if e.child.level != n.level-1 { 70 | t.Errorf("failed to preserve level order") 71 | } 72 | if e.child.parent != n { 73 | t.Errorf("failed to update parent pointer") 74 | } 75 | verify(t, e.child) 76 | } 77 | } 78 | 79 | func indexOf(objs []Spatial, obj Spatial) int { 80 | ind := -1 81 | for i, r := range objs { 82 | if r == obj { 83 | ind = i 84 | break 85 | } 86 | } 87 | return ind 88 | } 89 | 90 | var chooseLeafNodeTests = []struct { 91 | bb0, bb1, bb2 *Rect // leaf bounding boxes 92 | exp int // expected chosen leaf 93 | desc string 94 | level int 95 | }{ 96 | { 97 | mustRect(Point{1, 1, 1}, [Dim]float64{1, 1, 1}), 98 | mustRect(Point{-1, -1, -1}, [Dim]float64{0.5, 0.5, 0.5}), 99 | mustRect(Point{3, 4, -5}, [Dim]float64{2, 0.9, 8}), 100 | 1, 101 | "clear winner", 102 | 1, 103 | }, 104 | { 105 | mustRect(Point{-1, -1.5, -1}, [Dim]float64{0.5, 2.5025, 0.5}), 106 | mustRect(Point{0.5, 1, 0.5}, [Dim]float64{0.5, 0.815, 0.5}), 107 | mustRect(Point{3, 4, -5}, [Dim]float64{2, 0.9, 8}), 108 | 1, 109 | "leaves tie", 110 | 1, 111 | }, 112 | { 113 | mustRect(Point{-1, -1.5, -1}, [Dim]float64{0.5, 2.5025, 0.5}), 114 | mustRect(Point{0.5, 1, 0.5}, [Dim]float64{0.5, 0.815, 0.5}), 115 | mustRect(Point{-1, -2, -3}, [Dim]float64{2, 4, 6}), 116 | 2, 117 | "leaf contains obj", 118 | 1, 119 | }, 120 | } 121 | 122 | func TestChooseLeafNodeEmpty(t *testing.T) { 123 | rt := NewTree(5, 10) 124 | obj := Point{0, 0, 0}.ToRect(0.5) 125 | e := entry{obj, nil, obj} 126 | if leaf := rt.chooseNode(rt.root, e, 1); leaf != rt.root { 127 | t.Errorf("expected chooseLeaf of empty tree to return root") 128 | } 129 | } 130 | 131 | func TestChooseLeafNode(t *testing.T) { 132 | for _, test := range chooseLeafNodeTests { 133 | rt := Rtree{} 134 | rt.root = &node{} 135 | 136 | leaf0 := &node{rt.root, true, []entry{}, 1} 137 | entry0 := entry{test.bb0, leaf0, nil} 138 | 139 | leaf1 := &node{rt.root, true, []entry{}, 1} 140 | entry1 := entry{test.bb1, leaf1, nil} 141 | 142 | leaf2 := &node{rt.root, true, []entry{}, 1} 143 | entry2 := entry{test.bb2, leaf2, nil} 144 | 145 | rt.root.entries = []entry{entry0, entry1, entry2} 146 | 147 | obj := Point{0, 0, 0}.ToRect(0.5) 148 | e := entry{obj, nil, obj} 149 | 150 | expected := rt.root.entries[test.exp].child 151 | if leaf := rt.chooseNode(rt.root, e, 1); leaf != expected { 152 | t.Errorf("%s: expected %d", test.desc, test.exp) 153 | } 154 | } 155 | } 156 | 157 | func TestPickSeeds(t *testing.T) { 158 | entry1 := entry{bb: mustRect(Point{1, 1}, [Dim]float64{1, 1, 1})} 159 | entry2 := entry{bb: mustRect(Point{1, -1}, [Dim]float64{2, 1, 1})} 160 | entry3 := entry{bb: mustRect(Point{-1, -1}, [Dim]float64{1, 2, 1})} 161 | n := node{entries: []entry{entry1, entry2, entry3}} 162 | left, right := n.pickSeeds() 163 | if n.entries[left] != entry1 || n.entries[right] != entry3 { 164 | t.Errorf("expected entries %d, %d", 1, 3) 165 | } 166 | } 167 | 168 | func TestPickNext(t *testing.T) { 169 | leftEntry := entry{bb: mustRect(Point{1, 1}, [Dim]float64{1, 1, 1})} 170 | left := &node{entries: []entry{leftEntry}} 171 | 172 | rightEntry := entry{bb: mustRect(Point{-1, -1}, [Dim]float64{1, 2, 1})} 173 | right := &node{entries: []entry{rightEntry}} 174 | 175 | entry1 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1, 1})} 176 | entry2 := entry{bb: mustRect(Point{-2, -2}, [Dim]float64{1, 1, 1})} 177 | entry3 := entry{bb: mustRect(Point{1, 2}, [Dim]float64{1, 1, 1})} 178 | entries := []entry{entry1, entry2, entry3} 179 | 180 | chosen := pickNext(left, right, entries) 181 | if entries[chosen] != entry2 { 182 | t.Errorf("expected entry %d", 3) 183 | } 184 | } 185 | 186 | func TestSplit(t *testing.T) { 187 | entry1 := entry{bb: mustRect(Point{-3, -1}, [Dim]float64{2, 1, 1})} 188 | entry2 := entry{bb: mustRect(Point{1, 2}, [Dim]float64{1, 1, 1})} 189 | entry3 := entry{bb: mustRect(Point{-1, 0}, [Dim]float64{1, 1, 1})} 190 | entry4 := entry{bb: mustRect(Point{-3, -3}, [Dim]float64{1, 1, 1})} 191 | entry5 := entry{bb: mustRect(Point{1, -1}, [Dim]float64{2, 2, 1})} 192 | entries := []entry{entry1, entry2, entry3, entry4, entry5} 193 | n := &node{entries: entries} 194 | 195 | l, r := n.split(0) // left=entry2, right=entry4 196 | expLeft := mustRect(Point{1, -1}, [Dim]float64{2, 4, 1}) 197 | expRight := mustRect(Point{-3, -3}, [Dim]float64{3, 4, 1}) 198 | 199 | lbb := l.computeBoundingBox() 200 | rbb := r.computeBoundingBox() 201 | if lbb.p.dist(expLeft.p) >= EPS || lbb.q.dist(expLeft.q) >= EPS { 202 | t.Errorf("expected left.bb = %s, got %s", expLeft, lbb) 203 | } 204 | if rbb.p.dist(expRight.p) >= EPS || rbb.q.dist(expRight.q) >= EPS { 205 | t.Errorf("expected right.bb = %s, got %s", expRight, rbb) 206 | } 207 | } 208 | 209 | func TestSplitUnderflow(t *testing.T) { 210 | entry1 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1, 1})} 211 | entry2 := entry{bb: mustRect(Point{0, 1}, [Dim]float64{1, 1, 1})} 212 | entry3 := entry{bb: mustRect(Point{0, 2}, [Dim]float64{1, 1, 1})} 213 | entry4 := entry{bb: mustRect(Point{0, 3}, [Dim]float64{1, 1, 1})} 214 | entry5 := entry{bb: mustRect(Point{-50, -50}, [Dim]float64{1, 1, 1})} 215 | entries := []entry{entry1, entry2, entry3, entry4, entry5} 216 | n := &node{entries: entries} 217 | 218 | l, r := n.split(2) 219 | 220 | if len(l.entries) != 3 || len(r.entries) != 2 { 221 | t.Errorf("expected underflow assignment for right group") 222 | } 223 | } 224 | 225 | func TestAssignGroupLeastEnlargement(t *testing.T) { 226 | r00 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1, 1})} 227 | r01 := entry{bb: mustRect(Point{0, 1}, [Dim]float64{1, 1, 1})} 228 | r10 := entry{bb: mustRect(Point{1, 0}, [Dim]float64{1, 1, 1})} 229 | r11 := entry{bb: mustRect(Point{1, 1}, [Dim]float64{1, 1, 1})} 230 | r02 := entry{bb: mustRect(Point{0, 2}, [Dim]float64{1, 1, 1})} 231 | 232 | group1 := &node{entries: []entry{r00, r01}} 233 | group2 := &node{entries: []entry{r10, r11}} 234 | 235 | assignGroup(r02, group1, group2) 236 | if len(group1.entries) != 3 || len(group2.entries) != 2 { 237 | t.Errorf("expected r02 added to group 1") 238 | } 239 | } 240 | 241 | func TestAssignGroupSmallerArea(t *testing.T) { 242 | r00 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1, 1})} 243 | r01 := entry{bb: mustRect(Point{0, 1}, [Dim]float64{1, 1, 1})} 244 | r12 := entry{bb: mustRect(Point{1, 2}, [Dim]float64{1, 1, 1})} 245 | r02 := entry{bb: mustRect(Point{0, 2}, [Dim]float64{1, 1, 1})} 246 | 247 | group1 := &node{entries: []entry{r00, r01}} 248 | group2 := &node{entries: []entry{r12}} 249 | 250 | assignGroup(r02, group1, group2) 251 | if len(group2.entries) != 2 || len(group1.entries) != 2 { 252 | t.Errorf("expected r02 added to group 2") 253 | } 254 | } 255 | 256 | func TestAssignGroupFewerEntries(t *testing.T) { 257 | r0001 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 2, 1})} 258 | r12 := entry{bb: mustRect(Point{1, 2}, [Dim]float64{1, 1, 1})} 259 | r22 := entry{bb: mustRect(Point{2, 2}, [Dim]float64{1, 1, 1})} 260 | r02 := entry{bb: mustRect(Point{0, 2}, [Dim]float64{1, 1, 1})} 261 | 262 | group1 := &node{entries: []entry{r0001}} 263 | group2 := &node{entries: []entry{r12, r22}} 264 | 265 | assignGroup(r02, group1, group2) 266 | if len(group2.entries) != 2 || len(group1.entries) != 2 { 267 | t.Errorf("expected r02 added to group 2") 268 | } 269 | } 270 | 271 | func TestAdjustTreeNoPreviousSplit(t *testing.T) { 272 | rt := Rtree{root: &node{}} 273 | 274 | r00 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1, 1})} 275 | r01 := entry{bb: mustRect(Point{0, 1}, [Dim]float64{1, 1, 1})} 276 | r10 := entry{bb: mustRect(Point{1, 0}, [Dim]float64{1, 1, 1})} 277 | entries := []entry{r00, r01, r10} 278 | n := node{rt.root, false, entries, 1} 279 | rt.root.entries = []entry{entry{bb: Point{0, 0}.ToRect(0), child: &n}} 280 | 281 | rt.adjustTree(&n, nil) 282 | 283 | e := rt.root.entries[0] 284 | p, q := Point{0, 0, 0}, Point{2, 2, 1} 285 | if p.dist(e.bb.p) >= EPS || q.dist(e.bb.q) >= EPS { 286 | t.Errorf("Expected adjustTree to fit %v,%v,%v", r00.bb, r01.bb, r10.bb) 287 | } 288 | } 289 | 290 | func TestAdjustTreeNoSplit(t *testing.T) { 291 | rt := NewTree(3, 3) 292 | 293 | r00 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1})} 294 | r01 := entry{bb: mustRect(Point{0, 1}, [Dim]float64{1, 1})} 295 | left := node{rt.root, false, []entry{r00, r01}, 1} 296 | leftEntry := entry{bb: Point{0, 0}.ToRect(0), child: &left} 297 | 298 | r10 := entry{bb: mustRect(Point{1, 0}, [Dim]float64{1, 1})} 299 | r11 := entry{bb: mustRect(Point{1, 1}, [Dim]float64{1, 1})} 300 | right := node{rt.root, false, []entry{r10, r11}, 1} 301 | 302 | rt.root.entries = []entry{leftEntry} 303 | retl, retr := rt.adjustTree(&left, &right) 304 | 305 | if retl != rt.root || retr != nil { 306 | t.Errorf("Expected adjustTree didn't split the root") 307 | } 308 | 309 | entries := rt.root.entries 310 | if entries[0].child != &left || entries[1].child != &right { 311 | t.Errorf("Expected adjustTree keeps left and adds n in parent") 312 | } 313 | 314 | lbb, rbb := entries[0].bb, entries[1].bb 315 | if lbb.p.dist(Point{0, 0}) >= EPS || lbb.q.dist(Point{1, 2, 1}) >= EPS { 316 | t.Errorf("Expected adjustTree to adjust left bb") 317 | } 318 | if rbb.p.dist(Point{1, 0}) >= EPS || rbb.q.dist(Point{2, 2, 1}) >= EPS { 319 | t.Errorf("Expected adjustTree to adjust right bb") 320 | } 321 | } 322 | 323 | func TestAdjustTreeSplitParent(t *testing.T) { 324 | rt := NewTree(1, 1) 325 | 326 | r00 := entry{bb: mustRect(Point{0, 0}, [Dim]float64{1, 1})} 327 | r01 := entry{bb: mustRect(Point{0, 1}, [Dim]float64{1, 1})} 328 | left := node{rt.root, false, []entry{r00, r01}, 1} 329 | leftEntry := entry{bb: Point{0, 0}.ToRect(0), child: &left} 330 | 331 | r10 := entry{bb: mustRect(Point{1, 0}, [Dim]float64{1, 1})} 332 | r11 := entry{bb: mustRect(Point{1, 1}, [Dim]float64{1, 1})} 333 | right := node{rt.root, false, []entry{r10, r11}, 1} 334 | 335 | rt.root.entries = []entry{leftEntry} 336 | retl, retr := rt.adjustTree(&left, &right) 337 | 338 | if len(retl.entries) != 1 || len(retr.entries) != 1 { 339 | t.Errorf("Expected adjustTree distributed the entries") 340 | } 341 | 342 | lbb, rbb := retl.entries[0].bb, retr.entries[0].bb 343 | if lbb.p.dist(Point{0, 0}) >= EPS || lbb.q.dist(Point{1, 2, 1}) >= EPS { 344 | t.Errorf("Expected left split got left entry") 345 | } 346 | if rbb.p.dist(Point{1, 0}) >= EPS || rbb.q.dist(Point{2, 2, 1}) >= EPS { 347 | t.Errorf("Expected right split got right entry") 348 | } 349 | } 350 | 351 | func TestInsertRepeated(t *testing.T) { 352 | rt := NewTree(3, 5) 353 | thing := mustRect(Point{0, 0}, [Dim]float64{2, 1}) 354 | for i := 0; i < 6; i++ { 355 | rt.Insert(thing) 356 | } 357 | } 358 | 359 | func TestInsertNoSplit(t *testing.T) { 360 | rt := NewTree(3, 3) 361 | thing := mustRect(Point{0, 0}, [Dim]float64{2, 1}) 362 | rt.Insert(thing) 363 | 364 | if rt.Size() != 1 { 365 | t.Errorf("Insert failed to increase tree size") 366 | } 367 | 368 | if len(rt.root.entries) != 1 || rt.root.entries[0].obj.(*Rect) != thing { 369 | t.Errorf("Insert failed to insert thing into root entries") 370 | } 371 | } 372 | 373 | func TestInsertSplitRoot(t *testing.T) { 374 | rt := NewTree(3, 3) 375 | things := []*Rect{ 376 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 377 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 378 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 379 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 380 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 381 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 382 | } 383 | for _, thing := range things { 384 | rt.Insert(thing) 385 | } 386 | 387 | if rt.Size() != 6 { 388 | t.Errorf("Insert failed to insert") 389 | } 390 | 391 | if len(rt.root.entries) != 2 { 392 | t.Errorf("Insert failed to split") 393 | } 394 | 395 | left, right := rt.root.entries[0].child, rt.root.entries[1].child 396 | if len(left.entries) != 3 || len(right.entries) != 3 { 397 | t.Errorf("Insert failed to split evenly") 398 | } 399 | } 400 | 401 | func TestInsertSplit(t *testing.T) { 402 | rt := NewTree(3, 3) 403 | things := []*Rect{ 404 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 405 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 406 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 407 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 408 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 409 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 410 | mustRect(Point{10, 10}, [Dim]float64{2, 2}), 411 | } 412 | for _, thing := range things { 413 | rt.Insert(thing) 414 | } 415 | 416 | if rt.Size() != 7 { 417 | t.Errorf("Insert failed to insert") 418 | } 419 | 420 | if len(rt.root.entries) != 3 { 421 | t.Errorf("Insert failed to split") 422 | } 423 | 424 | a, b, c := rt.root.entries[0], rt.root.entries[1], rt.root.entries[2] 425 | if len(a.child.entries) != 3 || 426 | len(b.child.entries) != 3 || 427 | len(c.child.entries) != 1 { 428 | t.Errorf("Insert failed to split evenly") 429 | } 430 | } 431 | 432 | func TestInsertSplitSecondLevel(t *testing.T) { 433 | rt := NewTree(3, 3) 434 | things := []*Rect{ 435 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 436 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 437 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 438 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 439 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 440 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 441 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 442 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 443 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 444 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 445 | } 446 | for _, thing := range things { 447 | rt.Insert(thing) 448 | } 449 | 450 | if rt.Size() != 10 { 451 | t.Errorf("Insert failed to insert") 452 | } 453 | 454 | // should split root 455 | if len(rt.root.entries) != 2 { 456 | t.Errorf("Insert failed to split the root") 457 | } 458 | 459 | // split level + entries level + objs level 460 | if rt.Depth() != 3 { 461 | t.Errorf("Insert failed to adjust properly") 462 | } 463 | 464 | var checkParents func(n *node) 465 | checkParents = func(n *node) { 466 | if n.leaf { 467 | return 468 | } 469 | for _, e := range n.entries { 470 | if e.child.parent != n { 471 | t.Errorf("Insert failed to update parent pointers") 472 | } 473 | checkParents(e.child) 474 | } 475 | } 476 | checkParents(rt.root) 477 | } 478 | 479 | func TestFindLeaf(t *testing.T) { 480 | rt := NewTree(3, 3) 481 | things := []*Rect{ 482 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 483 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 484 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 485 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 486 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 487 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 488 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 489 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 490 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 491 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 492 | } 493 | for _, thing := range things { 494 | rt.Insert(thing) 495 | } 496 | verify(t, rt.root) 497 | for _, thing := range things { 498 | leaf := rt.findLeaf(rt.root, thing) 499 | if leaf == nil { 500 | printNode(rt.root, 0) 501 | t.Errorf("Unable to find leaf containing an entry after insertion!") 502 | } 503 | var found bool 504 | for _, other := range leaf.entries { 505 | if other.obj == thing { 506 | found = true 507 | break 508 | } 509 | } 510 | if !found { 511 | printNode(rt.root, 0) 512 | printNode(leaf, 0) 513 | t.Errorf("Entry %v not found in leaf node %v!", thing, leaf) 514 | } 515 | } 516 | } 517 | 518 | func TestFindLeafDoesNotExist(t *testing.T) { 519 | rt := NewTree(3, 3) 520 | things := []*Rect{ 521 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 522 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 523 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 524 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 525 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 526 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 527 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 528 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 529 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 530 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 531 | } 532 | for _, thing := range things { 533 | rt.Insert(thing) 534 | } 535 | 536 | obj := mustRect(Point{99, 99}, [Dim]float64{99, 99}) 537 | leaf := rt.findLeaf(rt.root, obj) 538 | if leaf != nil { 539 | t.Errorf("findLeaf failed to return nil for non-existent object") 540 | } 541 | } 542 | 543 | func TestCondenseTreeEliminate(t *testing.T) { 544 | rt := NewTree(3, 3) 545 | things := []*Rect{ 546 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 547 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 548 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 549 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 550 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 551 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 552 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 553 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 554 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 555 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 556 | } 557 | for _, thing := range things { 558 | rt.Insert(thing) 559 | } 560 | 561 | // delete entry 2 from parent entries 562 | parent := rt.root.entries[0].child.entries[1].child 563 | parent.entries = append(parent.entries[:2], parent.entries[3:]...) 564 | rt.condenseTree(parent) 565 | 566 | retrieved := []Spatial{} 567 | for obj := range items(rt.root) { 568 | retrieved = append(retrieved, obj) 569 | } 570 | 571 | if len(retrieved) != len(things)-1 { 572 | t.Errorf("condenseTree failed to reinsert upstream elements") 573 | } 574 | 575 | verify(t, rt.root) 576 | } 577 | 578 | func TestChooseNodeNonLeaf(t *testing.T) { 579 | rt := NewTree(3, 3) 580 | things := []*Rect{ 581 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 582 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 583 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 584 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 585 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 586 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 587 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 588 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 589 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 590 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 591 | } 592 | for _, thing := range things { 593 | rt.Insert(thing) 594 | } 595 | 596 | obj := mustRect(Point{0, 10}, [Dim]float64{1, 2}) 597 | e := entry{obj, nil, obj} 598 | n := rt.chooseNode(rt.root, e, 2) 599 | if n.level != 2 { 600 | t.Errorf("chooseNode failed to stop at desired level") 601 | } 602 | } 603 | 604 | func TestInsertNonLeaf(t *testing.T) { 605 | rt := NewTree(3, 3) 606 | things := []*Rect{ 607 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 608 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 609 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 610 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 611 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 612 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 613 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 614 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 615 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 616 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 617 | } 618 | for _, thing := range things { 619 | rt.Insert(thing) 620 | } 621 | 622 | obj := mustRect(Point{99, 99}, [Dim]float64{99, 99}) 623 | e := entry{obj, nil, obj} 624 | rt.insert(e, 2) 625 | 626 | expected := rt.root.entries[1].child 627 | if expected.entries[1].obj != obj { 628 | t.Errorf("insert failed to insert entry at correct level") 629 | } 630 | } 631 | 632 | func TestDeleteFlatten(t *testing.T) { 633 | rt := NewTree(3, 3) 634 | things := []*Rect{ 635 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 636 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 637 | } 638 | for _, thing := range things { 639 | rt.Insert(thing) 640 | } 641 | 642 | // make sure flattening didn't nuke the tree 643 | rt.Delete(things[0]) 644 | verify(t, rt.root) 645 | } 646 | 647 | func TestDelete(t *testing.T) { 648 | rt := NewTree(3, 3) 649 | things := []*Rect{ 650 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 651 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 652 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 653 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 654 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 655 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 656 | mustRect(Point{0, 6}, [Dim]float64{1, 2}), 657 | mustRect(Point{1, 6}, [Dim]float64{1, 2}), 658 | mustRect(Point{0, 8}, [Dim]float64{1, 2}), 659 | mustRect(Point{1, 8}, [Dim]float64{1, 2}), 660 | } 661 | for _, thing := range things { 662 | rt.Insert(thing) 663 | } 664 | 665 | verify(t, rt.root) 666 | 667 | things2 := []*Rect{} 668 | for len(things) > 0 { 669 | i := rand.Int() % len(things) 670 | things2 = append(things2, things[i]) 671 | things = append(things[:i], things[i+1:]...) 672 | } 673 | 674 | for i, thing := range things2 { 675 | ok := rt.Delete(thing) 676 | if !ok { 677 | t.Errorf("Thing %v was not found in tree during deletion", thing) 678 | return 679 | } 680 | 681 | if rt.Size() != len(things2)-i-1 { 682 | t.Errorf("Delete failed to remove %v", thing) 683 | return 684 | } 685 | verify(t, rt.root) 686 | } 687 | } 688 | 689 | func TestSearchIntersect(t *testing.T) { 690 | rt := NewTree(3, 3) 691 | things := []*Rect{ 692 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 693 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 694 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 695 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 696 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 697 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 698 | mustRect(Point{2, 6}, [Dim]float64{1, 2}), 699 | mustRect(Point{3, 6}, [Dim]float64{1, 2}), 700 | mustRect(Point{2, 8}, [Dim]float64{1, 2}), 701 | mustRect(Point{3, 8}, [Dim]float64{1, 2}), 702 | } 703 | for _, thing := range things { 704 | rt.Insert(thing) 705 | } 706 | 707 | bb := mustRect(Point{2, 1.5}, [Dim]float64{10, 5.5}) 708 | q := rt.SearchIntersect(bb) 709 | 710 | expected := []int{1, 2, 3, 4, 6, 7} 711 | if len(q) != len(expected) { 712 | t.Errorf("SearchIntersect failed to find all objects") 713 | } 714 | for _, ind := range expected { 715 | if indexOf(q, things[ind]) < 0 { 716 | t.Errorf("SearchIntersect failed to find things[%d]", ind) 717 | } 718 | } 719 | } 720 | 721 | func TestSearchIntersectNoResults(t *testing.T) { 722 | rt := NewTree(3, 3) 723 | things := []*Rect{ 724 | mustRect(Point{0, 0}, [Dim]float64{2, 1}), 725 | mustRect(Point{3, 1}, [Dim]float64{1, 2}), 726 | mustRect(Point{1, 2}, [Dim]float64{2, 2}), 727 | mustRect(Point{8, 6}, [Dim]float64{1, 1}), 728 | mustRect(Point{10, 3}, [Dim]float64{1, 2}), 729 | mustRect(Point{11, 7}, [Dim]float64{1, 1}), 730 | mustRect(Point{2, 6}, [Dim]float64{1, 2}), 731 | mustRect(Point{3, 6}, [Dim]float64{1, 2}), 732 | mustRect(Point{2, 8}, [Dim]float64{1, 2}), 733 | mustRect(Point{3, 8}, [Dim]float64{1, 2}), 734 | } 735 | for _, thing := range things { 736 | rt.Insert(thing) 737 | } 738 | 739 | bb := mustRect(Point{99, 99}, [Dim]float64{10, 5.5}) 740 | q := rt.SearchIntersect(bb) 741 | if len(q) != 0 { 742 | t.Errorf("SearchIntersect failed to return nil slice on failing query") 743 | } 744 | } 745 | 746 | func TestSortEntries(t *testing.T) { 747 | objs := []*Rect{ 748 | mustRect(Point{1, 1}, [Dim]float64{1, 1}), 749 | mustRect(Point{2, 2}, [Dim]float64{1, 1}), 750 | mustRect(Point{3, 3}, [Dim]float64{1, 1}), 751 | } 752 | entries := []entry{ 753 | entry{objs[2], nil, objs[2]}, 754 | entry{objs[1], nil, objs[1]}, 755 | entry{objs[0], nil, objs[0]}, 756 | } 757 | sorted, dists := sortEntries(Point{0, 0}, entries) 758 | if sorted[0] != entries[2] || sorted[1] != entries[1] || sorted[2] != entries[0] { 759 | t.Errorf("sortEntries failed") 760 | } 761 | if dists[0] != 2 || dists[1] != 8 || dists[2] != 18 { 762 | t.Errorf("sortEntries failed to calculate proper distances") 763 | } 764 | } 765 | 766 | func TestNearestNeighbor(t *testing.T) { 767 | rt := NewTree(3, 3) 768 | things := []*Rect{ 769 | mustRect(Point{1, 1}, [Dim]float64{1, 1}), 770 | mustRect(Point{1, 3}, [Dim]float64{1, 1}), 771 | mustRect(Point{3, 2}, [Dim]float64{1, 1}), 772 | mustRect(Point{-7, -7}, [Dim]float64{1, 1}), 773 | mustRect(Point{7, 7}, [Dim]float64{1, 1}), 774 | mustRect(Point{10, 2}, [Dim]float64{1, 1}), 775 | } 776 | for _, thing := range things { 777 | rt.Insert(thing) 778 | } 779 | 780 | obj1 := rt.NearestNeighbor(Point{0.5, 0.5}) 781 | obj2 := rt.NearestNeighbor(Point{1.5, 4.5}) 782 | obj3 := rt.NearestNeighbor(Point{5, 2.5}) 783 | obj4 := rt.NearestNeighbor(Point{3.5, 2.5}) 784 | 785 | if obj1 != things[0] || obj2 != things[1] || obj3 != things[2] || obj4 != things[2] { 786 | t.Errorf("NearestNeighbor failed") 787 | } 788 | } 789 | 790 | func TestNearestNeighbors(t *testing.T) { 791 | rt := NewTree(3, 3) 792 | things := []*Rect{ 793 | mustRect(Point{1, 1}, [Dim]float64{1, 1}), 794 | mustRect(Point{-7, -7}, [Dim]float64{1, 1}), 795 | mustRect(Point{1, 3}, [Dim]float64{1, 1}), 796 | mustRect(Point{7, 7}, [Dim]float64{1, 1}), 797 | mustRect(Point{10, 2}, [Dim]float64{1, 1}), 798 | mustRect(Point{3, 3}, [Dim]float64{1, 1}), 799 | } 800 | for _, thing := range things { 801 | rt.Insert(thing) 802 | } 803 | 804 | objs := rt.NearestNeighbors(3, Point{0.5, 0.5}) 805 | if objs[0] != things[0] || objs[1] != things[2] || objs[2] != things[5] { 806 | t.Errorf("NearestNeighbors failed") 807 | } 808 | } 809 | 810 | func BenchmarkSearchIntersect(b *testing.B) { 811 | b.StopTimer() 812 | rt := NewTree(3, 3) 813 | things := []*Rect{ 814 | mustRect(Point{0, 0, 0}, [Dim]float64{2, 1, 1}), 815 | mustRect(Point{3, 1, 0}, [Dim]float64{1, 2, 1}), 816 | mustRect(Point{1, 2, 0}, [Dim]float64{2, 2, 1}), 817 | mustRect(Point{8, 6, 0}, [Dim]float64{1, 1, 1}), 818 | mustRect(Point{10, 3, 0}, [Dim]float64{1, 2, 1}), 819 | mustRect(Point{11, 7, 0}, [Dim]float64{1, 1, 1}), 820 | mustRect(Point{2, 6, 0}, [Dim]float64{1, 2, 1}), 821 | mustRect(Point{3, 6, 0}, [Dim]float64{1, 2, 1}), 822 | mustRect(Point{2, 8, 0}, [Dim]float64{1, 2, 1}), 823 | mustRect(Point{3, 8, 0}, [Dim]float64{1, 2, 1}), 824 | } 825 | for _, thing := range things { 826 | rt.Insert(thing) 827 | } 828 | bb := mustRect(Point{2, 1.5, 0}, [Dim]float64{10, 5.5, 1}) 829 | b.StartTimer() 830 | for i := 0; i < b.N; i++ { 831 | rt.SearchIntersect(bb) 832 | } 833 | } 834 | 835 | func BenchmarkInsert(b *testing.B) { 836 | for i := 0; i < b.N; i++ { 837 | rt := NewTree(3, 3) 838 | things := []*Rect{ 839 | mustRect(Point{0, 0, 0}, [Dim]float64{2, 1, 1}), 840 | mustRect(Point{3, 1, 0}, [Dim]float64{1, 2, 1}), 841 | mustRect(Point{1, 2, 0}, [Dim]float64{2, 2, 1}), 842 | mustRect(Point{8, 6, 0}, [Dim]float64{1, 1, 1}), 843 | mustRect(Point{10, 3, 0}, [Dim]float64{1, 2, 1}), 844 | mustRect(Point{11, 7, 0}, [Dim]float64{1, 1, 1}), 845 | mustRect(Point{2, 6, 0}, [Dim]float64{1, 2, 1}), 846 | mustRect(Point{3, 6, 0}, [Dim]float64{1, 2, 1}), 847 | mustRect(Point{2, 8, 0}, [Dim]float64{1, 2, 1}), 848 | mustRect(Point{3, 8, 0}, [Dim]float64{1, 2, 1}), 849 | } 850 | for _, thing := range things { 851 | rt.Insert(thing) 852 | } 853 | } 854 | } 855 | --------------------------------------------------------------------------------