├── .gitignore ├── cities.png ├── go.mod ├── go.sum ├── qtree_test.go ├── LICENSE.md ├── README.md └── qtree.go /.gitignore: -------------------------------------------------------------------------------- 1 | *.svg -------------------------------------------------------------------------------- /cities.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tidwall/qtree/HEAD/cities.png -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/tidwall/qtree 2 | 3 | go 1.18 4 | 5 | require github.com/tidwall/geoindex v1.6.1 6 | 7 | require ( 8 | github.com/tidwall/cities v0.1.0 // indirect 9 | github.com/tidwall/lotsa v1.0.2 // indirect 10 | ) 11 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/tidwall/cities v0.1.0 h1:CVNkmMf7NEC9Bvokf5GoSsArHCKRMTgLuubRTHnH0mE= 2 | github.com/tidwall/cities v0.1.0/go.mod h1:lV/HDp2gCcRcHJWqgt6Di54GiDrTZwh1aG2ZUPNbqa4= 3 | github.com/tidwall/geoindex v1.6.1 h1:zGVJiAhHl5ZsPZHMM5HRpgv7Bqw3ZNxbnFRf3pxS1dk= 4 | github.com/tidwall/geoindex v1.6.1/go.mod h1:rvVVNEFfkJVWGUdEfU8QaoOg/9zFX0h9ofWzA60mz1I= 5 | github.com/tidwall/lotsa v1.0.2 h1:dNVBH5MErdaQ/xd9s769R31/n2dXavsQ0Yf4TMEHHw8= 6 | github.com/tidwall/lotsa v1.0.2/go.mod h1:X6NiU+4yHA3fE3Puvpnn1XMDrFZrE9JO2/w+UMuqgR8= 7 | -------------------------------------------------------------------------------- /qtree_test.go: -------------------------------------------------------------------------------- 1 | package qtree 2 | 3 | import ( 4 | "math/rand" 5 | "testing" 6 | "time" 7 | 8 | "github.com/tidwall/geoindex" 9 | ) 10 | 11 | func init() { 12 | seed := time.Now().UnixNano() 13 | println("seed:", seed) 14 | rand.Seed(seed) 15 | } 16 | 17 | func TestGeoIndex(t *testing.T) { 18 | t.Run("BenchVarious", func(t *testing.T) { 19 | geoindex.Tests.TestBenchVarious(t, &QTree{}, 1000000) 20 | }) 21 | t.Run("RandomRects", func(t *testing.T) { 22 | geoindex.Tests.TestRandomRects(t, &QTree{}, 10000) 23 | }) 24 | t.Run("RandomPoints", func(t *testing.T) { 25 | geoindex.Tests.TestRandomPoints(t, &QTree{}, 10000) 26 | }) 27 | t.Run("ZeroPoints", func(t *testing.T) { 28 | geoindex.Tests.TestZeroPoints(t, &QTree{}) 29 | }) 30 | t.Run("CitiesSVG", func(t *testing.T) { 31 | geoindex.Tests.TestCitiesSVG(t, &QTree{}) 32 | }) 33 | t.Run("RandomSVG", func(t *testing.T) { 34 | geoindex.Tests.TestRandomSVG(t, &QTree{}) 35 | }) 36 | } 37 | 38 | func BenchmarkRandomInsert(b *testing.B) { 39 | geoindex.Tests.BenchmarkRandomInsert(b, &QTree{}) 40 | } 41 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Copyright (c) 2020 Josh Baker 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in 11 | all copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 | THE SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # QTree 2 | 3 | [![GoDoc](https://godoc.org/github.com/tidwall/qtree?status.svg)](https://godoc.org/github.com/tidwall/qtree) 4 | 5 | This package provides an in-memory quadtree implementation for Go. 6 | 7 | Cities 8 | 9 | ## Usage 10 | 11 | ### Installing 12 | 13 | To start using QTree, install Go and run `go get`: 14 | 15 | ```sh 16 | $ go get -u github.com/tidwall/qtree 17 | ``` 18 | 19 | ### Basic operations 20 | 21 | ```go 22 | // create a QTree 23 | var tr qtree.QTree 24 | 25 | // insert a point 26 | tr.Insert([2]float64{-112.0078, 33.4373}, [2]float64{-112.0078, 33.4373}, "PHX") 27 | 28 | // insert a box 29 | tr.Insert([2]float64{10, 10}, [2]float64{20, 20}, "rect") 30 | 31 | // search 32 | tr.Search([2]float64{-112.1, 33.4}, [2]float64{-112.0, 33.5}, 33 | func(min, max [2]float64, value interface{}) bool { 34 | println(value.(string)) // prints "PHX" 35 | }, 36 | ) 37 | 38 | // delete 39 | tr.Delete([2]float64{-112.0078, 33.4373}, [2]float64{-112.0078, 33.4373}, "PHX") 40 | ``` 41 | 42 | ## Performance 43 | 44 | Pretty, pretty, pretty good. 45 | 46 | ## License 47 | 48 | `qtree` source code is available under the MIT License. 49 | -------------------------------------------------------------------------------- /qtree.go: -------------------------------------------------------------------------------- 1 | package qtree 2 | 3 | import ( 4 | "github.com/tidwall/geoindex" 5 | "github.com/tidwall/geoindex/child" 6 | ) 7 | 8 | const maxEntries = 32 9 | const minEntries = maxEntries * 40 / 100 10 | const maxLevel = 16 11 | 12 | // QTree ... 13 | type QTree struct { 14 | init bool 15 | min [2]float64 16 | max [2]float64 17 | root node 18 | } 19 | 20 | // New creates a new QuadTree 21 | func New(min, max [2]float64) *QTree { 22 | return &QTree{min: min, max: max, init: true} 23 | } 24 | 25 | var _ geoindex.Interface = &QTree{} 26 | 27 | // Insert an item into the structure 28 | func (tr *QTree) Insert(min [2]float64, max [2]float64, data interface{}) { 29 | if !tr.init { 30 | *tr = *New([2]float64{-180, -90}, [2]float64{180, 90}) 31 | } 32 | tr.root.insert(tr.min, tr.max, 1, entry{min, max, data}) 33 | } 34 | 35 | type entry struct { 36 | min, max [2]float64 37 | data interface{} 38 | } 39 | 40 | type node struct { 41 | count int 42 | entries []entry 43 | quads *[4]node 44 | } 45 | 46 | func (n *node) insert(min, max [2]float64, level int, e entry) { 47 | if n.quads == nil { 48 | if level != maxLevel && len(n.entries) == maxEntries { 49 | n.split(min, max, level) 50 | n.insert(min, max, level, e) 51 | return 52 | } 53 | n.entries = append(n.entries, e) 54 | } else { 55 | qmin, qmax, quad := choose(min, max, e) 56 | if quad == -1 { 57 | n.entries = append(n.entries, e) 58 | } else { 59 | n.quads[quad].insert(qmin, qmax, level+1, e) 60 | } 61 | } 62 | n.count++ 63 | } 64 | 65 | func (n *node) split(min, max [2]float64, level int) { 66 | entries := n.entries 67 | n.quads = new([4]node) 68 | n.entries = nil 69 | n.count = 0 70 | for _, e := range entries { 71 | n.insert(min, max, level, e) 72 | } 73 | } 74 | 75 | func (n *node) join() { 76 | n.entries = append(n.entries, n.quads[0].entries...) 77 | n.entries = append(n.entries, n.quads[1].entries...) 78 | n.entries = append(n.entries, n.quads[2].entries...) 79 | n.entries = append(n.entries, n.quads[3].entries...) 80 | n.quads = nil 81 | } 82 | 83 | func (n *node) delete(min, max [2]float64, e entry) bool { 84 | for i := 0; i < len(n.entries); i++ { 85 | if n.entries[i].data == e.data { 86 | n.entries[i] = n.entries[len(n.entries)-1] 87 | n.entries[len(n.entries)-1] = entry{} 88 | n.entries = n.entries[:len(n.entries)-1] 89 | n.count-- 90 | return true 91 | } 92 | } 93 | if n.quads != nil { 94 | qmin, qmax, quad := choose(min, max, e) 95 | if quad != -1 { 96 | if n.quads[quad].delete(qmin, qmax, e) { 97 | n.count-- 98 | if n.count < minEntries { 99 | n.join() 100 | } 101 | return true 102 | } 103 | } 104 | } 105 | return false 106 | } 107 | 108 | func (n *node) search( 109 | min, max [2]float64, tmin, tmax [2]float64, 110 | iter func(min, max [2]float64, data interface{}) bool, 111 | ) bool { 112 | for _, e := range n.entries { 113 | if intersects(tmin, tmax, e.min, e.max) { 114 | if !iter(e.min, e.max, e.data) { 115 | return false 116 | } 117 | } 118 | } 119 | if n.quads != nil { 120 | if n.quads[0].count > 0 { 121 | if qmin, qmax := q0(min, max); intersects(tmin, tmax, qmin, qmax) { 122 | if !n.quads[0].search(qmin, qmax, tmin, tmax, iter) { 123 | return false 124 | } 125 | } 126 | } 127 | if n.quads[1].count > 0 { 128 | if qmin, qmax := q1(min, max); intersects(tmin, tmax, qmin, qmax) { 129 | if !n.quads[1].search(qmin, qmax, tmin, tmax, iter) { 130 | return false 131 | } 132 | } 133 | } 134 | if n.quads[2].count > 0 { 135 | if qmin, qmax := q2(min, max); intersects(tmin, tmax, qmin, qmax) { 136 | if !n.quads[2].search(qmin, qmax, tmin, tmax, iter) { 137 | return false 138 | } 139 | } 140 | } 141 | if n.quads[3].count > 0 { 142 | if qmin, qmax := q3(min, max); intersects(tmin, tmax, qmin, qmax) { 143 | if !n.quads[3].search(qmin, qmax, tmin, tmax, iter) { 144 | return false 145 | } 146 | } 147 | } 148 | } 149 | return true 150 | } 151 | 152 | func (n *node) scan( 153 | iter func(min, max [2]float64, data interface{}) bool, 154 | ) bool { 155 | for _, e := range n.entries { 156 | if !iter(e.min, e.max, e.data) { 157 | return false 158 | } 159 | } 160 | if n.quads != nil { 161 | if !n.quads[0].scan(iter) { 162 | return false 163 | } 164 | if !n.quads[1].scan(iter) { 165 | return false 166 | } 167 | if !n.quads[2].scan(iter) { 168 | return false 169 | } 170 | if !n.quads[3].scan(iter) { 171 | return false 172 | } 173 | } 174 | return true 175 | } 176 | 177 | // Delete an item from the structure 178 | func (tr *QTree) Delete(min, max [2]float64, data interface{}) { 179 | if !tr.init { 180 | return 181 | } 182 | tr.root.delete(tr.min, tr.max, entry{min, max, data}) 183 | } 184 | 185 | // Replace an item in the structure. This is effectively just a Delete 186 | // followed by an Insert. But for some structures it may be possible to 187 | // optimize the operation to avoid multiple passes 188 | func (tr *QTree) Replace( 189 | oldMin, oldMax [2]float64, oldData interface{}, 190 | newMin, newMax [2]float64, newData interface{}, 191 | ) { 192 | tr.Delete(oldMin, oldMax, oldData) 193 | tr.Insert(newMin, newMax, newData) 194 | } 195 | 196 | // Search the structure for items that intersects the rect param 197 | func (tr *QTree) Search( 198 | min, max [2]float64, 199 | iter func(min, max [2]float64, data interface{}) bool, 200 | ) { 201 | if !tr.init { 202 | return 203 | } 204 | tr.root.search(tr.min, tr.max, min, max, iter) 205 | } 206 | 207 | // Scan iterates through all data in tree in no specified order. 208 | func (tr *QTree) Scan(iter func(min, max [2]float64, data interface{}) bool) { 209 | tr.root.scan(iter) 210 | } 211 | 212 | // Len returns the number of items in tree 213 | func (tr *QTree) Len() int { 214 | return tr.root.count 215 | } 216 | 217 | // Bounds returns the minimum bounding box 218 | func (tr *QTree) Bounds() (min, max [2]float64) { 219 | if !tr.init { 220 | return min, max 221 | } 222 | minX, _ := tr.root.min(tr.min, tr.max, 0, 0, false) 223 | minY, _ := tr.root.min(tr.min, tr.max, 1, 0, false) 224 | maxX, _ := tr.root.max(tr.min, tr.max, 0, 0, false) 225 | maxY, _ := tr.root.max(tr.min, tr.max, 1, 0, false) 226 | return [2]float64{minX, minY}, [2]float64{maxX, maxY} 227 | } 228 | 229 | type rnode struct { 230 | min, max [2]float64 231 | node *node 232 | } 233 | 234 | // Children returns all children for parent node. If parent node is nil 235 | // then the root nodes should be returned. 236 | // The reuse buffer is an empty length slice that can optionally be used 237 | // to avoid extra allocations. 238 | func (tr *QTree) Children(parent interface{}, reuse []child.Child) ( 239 | children []child.Child, 240 | ) { 241 | children = reuse[:0] 242 | if parent == nil { 243 | if tr.Len() > 0 { 244 | // fill with the root 245 | children = append(children, child.Child{ 246 | Min: tr.min, 247 | Max: tr.max, 248 | Data: &rnode{tr.min, tr.max, &tr.root}, 249 | Item: false, 250 | }) 251 | } 252 | } else { 253 | // fill with child items 254 | switch n := parent.(type) { 255 | case *rnode: 256 | if n.node.quads == nil { 257 | // leaf 258 | for _, e := range n.node.entries { 259 | children = append(children, child.Child{ 260 | Min: e.min, 261 | Max: e.max, 262 | Data: e.data, 263 | Item: true, 264 | }) 265 | } 266 | return children 267 | } 268 | 269 | if len(n.node.entries) > 0 { 270 | children = append(children, child.Child{ 271 | Min: n.min, 272 | Max: n.max, 273 | Data: n.node.entries, 274 | Item: false, 275 | }) 276 | } 277 | if n.node.quads[0].count > 0 { 278 | qmin, qmax := q0(n.min, n.max) 279 | children = append(children, child.Child{ 280 | Min: qmin, 281 | Max: qmax, 282 | Data: &rnode{qmin, qmax, &n.node.quads[0]}, 283 | Item: false, 284 | }) 285 | } 286 | if n.node.quads[1].count > 0 { 287 | qmin, qmax := q1(n.min, n.max) 288 | children = append(children, child.Child{ 289 | Min: qmin, 290 | Max: qmax, 291 | Data: &rnode{qmin, qmax, &n.node.quads[1]}, 292 | Item: false, 293 | }) 294 | } 295 | if n.node.quads[2].count > 0 { 296 | qmin, qmax := q2(n.min, n.max) 297 | children = append(children, child.Child{ 298 | Min: qmin, 299 | Max: qmax, 300 | Data: &rnode{qmin, qmax, &n.node.quads[2]}, 301 | Item: false, 302 | }) 303 | } 304 | if n.node.quads[3].count > 0 { 305 | qmin, qmax := q3(n.min, n.max) 306 | children = append(children, child.Child{ 307 | Min: qmin, 308 | Max: qmax, 309 | Data: &rnode{qmin, qmax, &n.node.quads[3]}, 310 | Item: false, 311 | }) 312 | } 313 | case []entry: 314 | for _, e := range n { 315 | children = append(children, child.Child{ 316 | Min: e.min, 317 | Max: e.max, 318 | Data: e.data, 319 | Item: true, 320 | }) 321 | } 322 | } 323 | } 324 | return children 325 | } 326 | 327 | func intersects(amin, amax, bmin, bmax [2]float64) bool { 328 | if bmin[0] > amax[0] || bmax[0] < amin[0] { 329 | return false 330 | } 331 | if bmin[1] > amax[1] || bmax[1] < amin[1] { 332 | return false 333 | } 334 | return true 335 | } 336 | 337 | func contains(amin, amax, bmin, bmax [2]float64) bool { 338 | return bmin[0] >= amin[0] && bmax[0] <= amax[0] && 339 | bmin[1] >= amin[1] && bmax[1] <= amax[1] 340 | } 341 | 342 | func q0min(min, max [2]float64) [2]float64 { 343 | return [2]float64{min[0], min[1]} 344 | } 345 | func q0max(min, max [2]float64) [2]float64 { 346 | return [2]float64{(max[0] + min[0]) / 2, (max[1] + min[1]) / 2} 347 | } 348 | func q0(min, max [2]float64) ([2]float64, [2]float64) { 349 | return q0min(min, max), q0max(min, max) 350 | } 351 | 352 | func q1min(min, max [2]float64) [2]float64 { 353 | return [2]float64{(max[0] + min[0]) / 2, min[1]} 354 | } 355 | func q1max(min, max [2]float64) [2]float64 { 356 | return [2]float64{max[0], (max[1] + min[1]) / 2} 357 | } 358 | func q1(min, max [2]float64) ([2]float64, [2]float64) { 359 | return q1min(min, max), q1max(min, max) 360 | } 361 | 362 | func q2min(min, max [2]float64) [2]float64 { 363 | return [2]float64{min[0], (max[1] + min[1]) / 2} 364 | } 365 | func q2max(min, max [2]float64) [2]float64 { 366 | return [2]float64{(max[0] + min[0]) / 2, max[1]} 367 | } 368 | func q2(min, max [2]float64) ([2]float64, [2]float64) { 369 | return q2min(min, max), q2max(min, max) 370 | } 371 | 372 | func q3min(min, max [2]float64) [2]float64 { 373 | return [2]float64{(max[0] + min[0]) / 2, (max[1] + min[1]) / 2} 374 | } 375 | func q3max(min, max [2]float64) [2]float64 { 376 | return [2]float64{max[0], max[1]} 377 | } 378 | func q3(min, max [2]float64) ([2]float64, [2]float64) { 379 | return q3min(min, max), q3max(min, max) 380 | } 381 | 382 | func choose(min, max [2]float64, e entry) (qmin, qmax [2]float64, quad int) { 383 | if qmin, qmax = q0(min, max); contains(qmin, qmax, e.min, e.max) { 384 | return qmin, qmax, 0 385 | } 386 | if qmin, qmax = q1(min, max); contains(qmin, qmax, e.min, e.max) { 387 | return qmin, qmax, 1 388 | } 389 | if qmin, qmax = q2(min, max); contains(qmin, qmax, e.min, e.max) { 390 | return qmin, qmax, 2 391 | } 392 | if qmin, qmax = q3(min, max); contains(qmin, qmax, e.min, e.max) { 393 | return qmin, qmax, 3 394 | } 395 | return qmin, qmax, -1 396 | } 397 | 398 | func (n *node) min(min, max [2]float64, axis int, v float64, set bool) ( 399 | float64, bool, 400 | ) { 401 | for _, e := range n.entries { 402 | if !set || e.min[axis] < v { 403 | v, set = e.min[axis], true 404 | } 405 | } 406 | if (set && v < min[axis]) || n.quads == nil { 407 | return v, set 408 | } 409 | if axis == 0 { 410 | qmin, qmax := q0(min, max) 411 | qv, qset := n.quads[0].min(qmin, qmax, axis, v, set) 412 | if qset && (!set || qv < v) { 413 | v, set = qv, true 414 | } 415 | qmin, qmax = q2(min, max) 416 | qv, qset = n.quads[2].min(qmin, qmax, axis, v, set) 417 | if qset && (!set || qv < v) { 418 | v, set = qv, true 419 | } 420 | } 421 | return v, set 422 | } 423 | 424 | func (n *node) max(min, max [2]float64, axis int, v float64, set bool) ( 425 | float64, bool, 426 | ) { 427 | for _, e := range n.entries { 428 | if !set || e.max[axis] > v { 429 | v, set = e.max[axis], true 430 | } 431 | } 432 | if (set && v > max[axis]) || n.quads == nil { 433 | return v, set 434 | } 435 | if axis == 0 { 436 | qmin, qmax := q0(min, max) 437 | qv, qset := n.quads[0].max(qmin, qmax, axis, v, set) 438 | if qset && (!set || qv > v) { 439 | v, set = qv, true 440 | } 441 | qmin, qmax = q2(min, max) 442 | qv, qset = n.quads[2].max(qmin, qmax, axis, v, set) 443 | if qset && (!set || qv > v) { 444 | v, set = qv, true 445 | } 446 | } 447 | return v, set 448 | } 449 | --------------------------------------------------------------------------------