rtree.go 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674
  1. package base
  2. import (
  3. "math"
  4. "unsafe"
  5. )
  6. // precalculate infinity
  7. var mathInfNeg = math.Inf(-1)
  8. var mathInfPos = math.Inf(+1)
  9. type treeNode struct {
  10. min, max []float64
  11. children []*treeNode
  12. count int
  13. height int
  14. leaf bool
  15. }
  16. func (node *treeNode) unsafeItem() *treeItem {
  17. return (*treeItem)(unsafe.Pointer(node))
  18. }
  19. func (tr *RTree) createNode(children []*treeNode) *treeNode {
  20. n := &treeNode{
  21. height: 1,
  22. leaf: true,
  23. children: make([]*treeNode, tr.maxEntries+1),
  24. }
  25. if len(children) > 0 {
  26. n.count = len(children)
  27. copy(n.children[:n.count], children)
  28. }
  29. n.min = make([]float64, tr.dims)
  30. n.max = make([]float64, tr.dims)
  31. for i := 0; i < tr.dims; i++ {
  32. n.min[i] = mathInfPos
  33. n.max[i] = mathInfNeg
  34. }
  35. return n
  36. }
  37. func (node *treeNode) extend(b *treeNode) {
  38. for i := 0; i < len(node.min); i++ {
  39. if b.min[i] < node.min[i] {
  40. node.min[i] = b.min[i]
  41. }
  42. if b.max[i] > node.max[i] {
  43. node.max[i] = b.max[i]
  44. }
  45. }
  46. }
  47. func (node *treeNode) area() float64 {
  48. area := node.max[0] - node.min[0]
  49. for i := 1; i < len(node.min); i++ {
  50. area *= node.max[i] - node.min[i]
  51. }
  52. return area
  53. }
  54. func (node *treeNode) enlargedAreaAxis(b *treeNode, axis int) float64 {
  55. var max, min float64
  56. if b.max[axis] > node.max[axis] {
  57. max = b.max[axis]
  58. } else {
  59. max = node.max[axis]
  60. }
  61. if b.min[axis] < node.min[axis] {
  62. min = b.min[axis]
  63. } else {
  64. min = node.min[axis]
  65. }
  66. return max - min
  67. }
  68. func (node *treeNode) enlargedArea(b *treeNode) float64 {
  69. area := node.enlargedAreaAxis(b, 0)
  70. for i := 1; i < len(node.min); i++ {
  71. area *= node.enlargedAreaAxis(b, i)
  72. }
  73. return area
  74. }
  75. func (node *treeNode) intersectionAreaAxis(b *treeNode, axis int) float64 {
  76. var max, min float64
  77. if node.max[axis] < b.max[axis] {
  78. max = node.max[axis]
  79. } else {
  80. max = b.max[axis]
  81. }
  82. if node.min[axis] > b.min[axis] {
  83. min = node.min[axis]
  84. } else {
  85. min = b.min[axis]
  86. }
  87. if max > min {
  88. return max - min
  89. }
  90. return 0
  91. }
  92. func (node *treeNode) intersectionArea(b *treeNode) float64 {
  93. area := node.intersectionAreaAxis(b, 0)
  94. for i := 1; i < len(node.min); i++ {
  95. area *= node.intersectionAreaAxis(b, i)
  96. }
  97. return area
  98. }
  99. func (node *treeNode) margin() float64 {
  100. margin := node.max[0] - node.min[0]
  101. for i := 1; i < len(node.min); i++ {
  102. margin += node.max[i] - node.min[i]
  103. }
  104. return margin
  105. }
  106. type result int
  107. const (
  108. not result = 0
  109. intersects result = 1
  110. contains result = 2
  111. )
  112. func (node *treeNode) overlaps(b *treeNode) result {
  113. for i := 0; i < len(node.min); i++ {
  114. if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
  115. return not
  116. }
  117. if node.min[i] > b.min[i] || b.max[i] > node.max[i] {
  118. i++
  119. for ; i < len(node.min); i++ {
  120. if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
  121. return not
  122. }
  123. }
  124. return intersects
  125. }
  126. }
  127. return contains
  128. }
  129. func (node *treeNode) intersects(b *treeNode) bool {
  130. for i := 0; i < len(node.min); i++ {
  131. if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
  132. return false
  133. }
  134. }
  135. return true
  136. }
  137. func (node *treeNode) findItem(item interface{}) int {
  138. for i := 0; i < node.count; i++ {
  139. if node.children[i].unsafeItem().item == item {
  140. return i
  141. }
  142. }
  143. return -1
  144. }
  145. func (node *treeNode) contains(b *treeNode) bool {
  146. for i := 0; i < len(node.min); i++ {
  147. if node.min[i] > b.min[i] || b.max[i] > node.max[i] {
  148. return false
  149. }
  150. }
  151. return true
  152. }
  153. func (node *treeNode) childCount() int {
  154. if node.leaf {
  155. return node.count
  156. }
  157. var n int
  158. for i := 0; i < node.count; i++ {
  159. n += node.children[i].childCount()
  160. }
  161. return n
  162. }
  163. type treeItem struct {
  164. min, max []float64
  165. item interface{}
  166. }
  167. //go:nocheckptr
  168. func (item *treeItem) unsafeNode() *treeNode {
  169. return (*treeNode)(unsafe.Pointer(item))
  170. }
  171. // RTree is an R-tree
  172. type RTree struct {
  173. dims int
  174. maxEntries int
  175. minEntries int
  176. data *treeNode // root node
  177. // resusable fields, these help performance of common mutable operations.
  178. reuse struct {
  179. path []*treeNode // for reinsertion path
  180. indexes []int // for remove function
  181. stack []int // for bulk loading
  182. }
  183. }
  184. // New creates a new R-tree
  185. func New(dims, maxEntries int) *RTree {
  186. if dims <= 0 {
  187. panic("invalid dimensions")
  188. }
  189. tr := &RTree{}
  190. tr.dims = dims
  191. tr.maxEntries = int(math.Max(4, float64(maxEntries)))
  192. tr.minEntries = int(math.Max(2, math.Ceil(float64(tr.maxEntries)*0.4)))
  193. tr.data = tr.createNode(nil)
  194. return tr
  195. }
  196. // Insert inserts an item
  197. func (tr *RTree) Insert(min, max []float64, item interface{}) {
  198. if len(min) != tr.dims || len(max) != tr.dims {
  199. panic("invalid dimensions")
  200. }
  201. if item == nil {
  202. panic("nil item")
  203. }
  204. bbox := treeNode{min: min, max: max}
  205. tr.insert(&bbox, item, tr.data.height-1, false)
  206. }
  207. func (tr *RTree) insert(bbox *treeNode, item interface{}, level int, isNode bool) {
  208. tr.reuse.path = tr.reuse.path[:0]
  209. node, insertPath := tr.chooseSubtree(bbox, tr.data, level, tr.reuse.path)
  210. if item == nil {
  211. // item is only nil when bulk loading a node
  212. if node.leaf {
  213. panic("loading node into leaf")
  214. }
  215. node.children[node.count] = bbox
  216. node.count++
  217. } else {
  218. ti := &treeItem{min: bbox.min, max: bbox.max, item: item}
  219. node.children[node.count] = ti.unsafeNode()
  220. node.count++
  221. }
  222. node.extend(bbox)
  223. for level >= 0 {
  224. if insertPath[level].count > tr.maxEntries {
  225. insertPath = tr.split(insertPath, level)
  226. level--
  227. } else {
  228. break
  229. }
  230. }
  231. tr.adjustParentBBoxes(bbox, insertPath, level)
  232. tr.reuse.path = insertPath
  233. }
  234. func (tr *RTree) adjustParentBBoxes(bbox *treeNode, path []*treeNode, level int) {
  235. // adjust bboxes along the given tree path
  236. for i := level; i >= 0; i-- {
  237. path[i].extend(bbox)
  238. }
  239. }
  240. func (tr *RTree) chooseSubtree(bbox, node *treeNode, level int, path []*treeNode) (*treeNode, []*treeNode) {
  241. var targetNode *treeNode
  242. var area, enlargement, minArea, minEnlargement float64
  243. for {
  244. path = append(path, node)
  245. if node.leaf || len(path)-1 == level {
  246. break
  247. }
  248. minEnlargement = mathInfPos
  249. minArea = minEnlargement
  250. for i := 0; i < node.count; i++ {
  251. child := node.children[i]
  252. area = child.area()
  253. enlargement = bbox.enlargedArea(child) - area
  254. if enlargement < minEnlargement {
  255. minEnlargement = enlargement
  256. if area < minArea {
  257. minArea = area
  258. }
  259. targetNode = child
  260. } else if enlargement == minEnlargement {
  261. if area < minArea {
  262. minArea = area
  263. targetNode = child
  264. }
  265. }
  266. }
  267. if targetNode != nil {
  268. node = targetNode
  269. } else if node.count > 0 {
  270. node = (*treeNode)(node.children[0])
  271. } else {
  272. node = nil
  273. }
  274. }
  275. return node, path
  276. }
  277. func (tr *RTree) split(insertPath []*treeNode, level int) []*treeNode {
  278. var node = insertPath[level]
  279. var M = node.count
  280. var m = tr.minEntries
  281. tr.chooseSplitAxis(node, m, M)
  282. splitIndex := tr.chooseSplitIndex(node, m, M)
  283. spliced := make([]*treeNode, node.count-splitIndex)
  284. copy(spliced, node.children[splitIndex:])
  285. node.count = splitIndex
  286. newNode := tr.createNode(spliced)
  287. newNode.height = node.height
  288. newNode.leaf = node.leaf
  289. tr.calcBBox(node)
  290. tr.calcBBox(newNode)
  291. if level != 0 {
  292. insertPath[level-1].children[insertPath[level-1].count] = newNode
  293. insertPath[level-1].count++
  294. } else {
  295. tr.splitRoot(node, newNode)
  296. }
  297. return insertPath
  298. }
  299. func (tr *RTree) chooseSplitIndex(node *treeNode, m, M int) int {
  300. var i int
  301. var bbox1, bbox2 *treeNode
  302. var overlap, area, minOverlap, minArea float64
  303. var index int
  304. minArea = mathInfPos
  305. minOverlap = minArea
  306. for i = m; i <= M-m; i++ {
  307. bbox1 = tr.distBBox(node, 0, i, nil)
  308. bbox2 = tr.distBBox(node, i, M, nil)
  309. overlap = bbox1.intersectionArea(bbox2)
  310. area = bbox1.area() + bbox2.area()
  311. // choose distribution with minimum overlap
  312. if overlap < minOverlap {
  313. minOverlap = overlap
  314. index = i
  315. if area < minArea {
  316. minArea = area
  317. }
  318. } else if overlap == minOverlap {
  319. // otherwise choose distribution with minimum area
  320. if area < minArea {
  321. minArea = area
  322. index = i
  323. }
  324. }
  325. }
  326. return index
  327. }
  328. func (tr *RTree) calcBBox(node *treeNode) {
  329. tr.distBBox(node, 0, node.count, node)
  330. }
  331. func (tr *RTree) chooseSplitAxis(node *treeNode, m, M int) {
  332. minMargin := tr.allDistMargin(node, m, M, 0)
  333. var minAxis int
  334. for axis := 1; axis < tr.dims; axis++ {
  335. margin := tr.allDistMargin(node, m, M, axis)
  336. if margin < minMargin {
  337. minMargin = margin
  338. minAxis = axis
  339. }
  340. }
  341. if minAxis < tr.dims {
  342. tr.sortNodes(node, minAxis)
  343. }
  344. }
  345. func (tr *RTree) splitRoot(node, newNode *treeNode) {
  346. tr.data = tr.createNode([]*treeNode{node, newNode})
  347. tr.data.height = node.height + 1
  348. tr.data.leaf = false
  349. tr.calcBBox(tr.data)
  350. }
  351. func (tr *RTree) distBBox(node *treeNode, k, p int, destNode *treeNode) *treeNode {
  352. if destNode == nil {
  353. destNode = tr.createNode(nil)
  354. } else {
  355. for i := 0; i < tr.dims; i++ {
  356. destNode.min[i] = mathInfPos
  357. destNode.max[i] = mathInfNeg
  358. }
  359. }
  360. for i := k; i < p; i++ {
  361. if node.leaf {
  362. destNode.extend(node.children[i])
  363. } else {
  364. destNode.extend((*treeNode)(node.children[i]))
  365. }
  366. }
  367. return destNode
  368. }
  369. func (tr *RTree) allDistMargin(node *treeNode, m, M int, axis int) float64 {
  370. tr.sortNodes(node, axis)
  371. var leftBBox = tr.distBBox(node, 0, m, nil)
  372. var rightBBox = tr.distBBox(node, M-m, M, nil)
  373. var margin = leftBBox.margin() + rightBBox.margin()
  374. var i int
  375. if node.leaf {
  376. for i = m; i < M-m; i++ {
  377. leftBBox.extend(node.children[i])
  378. margin += leftBBox.margin()
  379. }
  380. for i = M - m - 1; i >= m; i-- {
  381. leftBBox.extend(node.children[i])
  382. margin += rightBBox.margin()
  383. }
  384. } else {
  385. for i = m; i < M-m; i++ {
  386. child := (*treeNode)(node.children[i])
  387. leftBBox.extend(child)
  388. margin += leftBBox.margin()
  389. }
  390. for i = M - m - 1; i >= m; i-- {
  391. child := (*treeNode)(node.children[i])
  392. leftBBox.extend(child)
  393. margin += rightBBox.margin()
  394. }
  395. }
  396. return margin
  397. }
  398. func (tr *RTree) sortNodes(node *treeNode, axis int) {
  399. sortByAxis(node.children[:node.count], axis)
  400. }
  401. func sortByAxis(items []*treeNode, axis int) {
  402. if len(items) < 2 {
  403. return
  404. }
  405. left, right := 0, len(items)-1
  406. pivotIndex := len(items) / 2
  407. items[pivotIndex], items[right] = items[right], items[pivotIndex]
  408. for i := range items {
  409. if items[i].min[axis] < items[right].min[axis] {
  410. items[i], items[left] = items[left], items[i]
  411. left++
  412. }
  413. }
  414. items[left], items[right] = items[right], items[left]
  415. sortByAxis(items[:left], axis)
  416. sortByAxis(items[left+1:], axis)
  417. }
  418. // Search searches the tree for items in the input rectangle
  419. func (tr *RTree) Search(min, max []float64, iter func(item interface{}) bool) bool {
  420. bbox := &treeNode{min: min, max: max}
  421. if !tr.data.intersects(bbox) {
  422. return true
  423. }
  424. return tr.search(tr.data, bbox, iter)
  425. }
  426. func (tr *RTree) search(node, bbox *treeNode, iter func(item interface{}) bool) bool {
  427. if node.leaf {
  428. for i := 0; i < node.count; i++ {
  429. if bbox.intersects(node.children[i]) {
  430. if !iter(node.children[i].unsafeItem().item) {
  431. return false
  432. }
  433. }
  434. }
  435. } else {
  436. for i := 0; i < node.count; i++ {
  437. r := bbox.overlaps(node.children[i])
  438. if r == intersects {
  439. if !tr.search(node.children[i], bbox, iter) {
  440. return false
  441. }
  442. } else if r == contains {
  443. if !scan(node.children[i], iter) {
  444. return false
  445. }
  446. }
  447. }
  448. }
  449. return true
  450. }
  451. func (tr *RTree) IsEmpty() bool {
  452. empty := true
  453. tr.Scan(func(item interface{}) bool {
  454. empty = false
  455. return false
  456. })
  457. return empty
  458. }
  459. // Remove removes an item from the R-tree.
  460. func (tr *RTree) Remove(min, max []float64, item interface{}) {
  461. bbox := &treeNode{min: min, max: max}
  462. tr.remove(bbox, item)
  463. }
  464. func (tr *RTree) remove(bbox *treeNode, item interface{}) {
  465. path := tr.reuse.path[:0]
  466. indexes := tr.reuse.indexes[:0]
  467. var node = tr.data
  468. var i int
  469. var parent *treeNode
  470. var index int
  471. var goingUp bool
  472. for node != nil || len(path) != 0 {
  473. if node == nil {
  474. node = path[len(path)-1]
  475. path = path[:len(path)-1]
  476. if len(path) == 0 {
  477. parent = nil
  478. } else {
  479. parent = path[len(path)-1]
  480. }
  481. i = indexes[len(indexes)-1]
  482. indexes = indexes[:len(indexes)-1]
  483. goingUp = true
  484. }
  485. if node.leaf {
  486. index = node.findItem(item)
  487. if index != -1 {
  488. // item found, remove the item and condense tree upwards
  489. copy(node.children[index:], node.children[index+1:])
  490. node.children[node.count-1] = nil
  491. node.count--
  492. path = append(path, node)
  493. tr.condense(path)
  494. goto done
  495. }
  496. }
  497. if !goingUp && !node.leaf && node.contains(bbox) { // go down
  498. path = append(path, node)
  499. indexes = append(indexes, i)
  500. i = 0
  501. parent = node
  502. node = (*treeNode)(node.children[0])
  503. } else if parent != nil { // go right
  504. i++
  505. if i == parent.count {
  506. node = nil
  507. } else {
  508. node = (*treeNode)(parent.children[i])
  509. }
  510. goingUp = false
  511. } else {
  512. node = nil
  513. }
  514. }
  515. done:
  516. tr.reuse.path = path
  517. tr.reuse.indexes = indexes
  518. return
  519. }
  520. func (tr *RTree) condense(path []*treeNode) {
  521. // go through the path, removing empty nodes and updating bboxes
  522. var siblings []*treeNode
  523. for i := len(path) - 1; i >= 0; i-- {
  524. if path[i].count == 0 {
  525. if i > 0 {
  526. siblings = path[i-1].children[:path[i-1].count]
  527. index := -1
  528. for j := 0; j < len(siblings); j++ {
  529. if siblings[j] == path[i] {
  530. index = j
  531. break
  532. }
  533. }
  534. copy(siblings[index:], siblings[index+1:])
  535. siblings[len(siblings)-1] = nil
  536. path[i-1].count--
  537. //siblings = siblings[:len(siblings)-1]
  538. //path[i-1].children = siblings
  539. } else {
  540. tr.data = tr.createNode(nil) // clear tree
  541. }
  542. } else {
  543. tr.calcBBox(path[i])
  544. }
  545. }
  546. }
  547. // Count returns the number of items in the R-tree.
  548. func (tr *RTree) Count() int {
  549. return tr.data.childCount()
  550. }
  551. // Traverse iterates over the entire R-tree and includes all nodes and items.
  552. func (tr *RTree) Traverse(iter func(min, max []float64, level int, item interface{}) bool) bool {
  553. return tr.traverse(tr.data, iter)
  554. }
  555. func (tr *RTree) traverse(node *treeNode, iter func(min, max []float64, level int, item interface{}) bool) bool {
  556. if !iter(node.min, node.max, int(node.height), nil) {
  557. return false
  558. }
  559. if node.leaf {
  560. for i := 0; i < node.count; i++ {
  561. child := node.children[i]
  562. if !iter(child.min, child.max, 0, child.unsafeItem().item) {
  563. return false
  564. }
  565. }
  566. } else {
  567. for i := 0; i < node.count; i++ {
  568. child := node.children[i]
  569. if !tr.traverse(child, iter) {
  570. return false
  571. }
  572. }
  573. }
  574. return true
  575. }
  576. // Scan iterates over the entire R-tree
  577. func (tr *RTree) Scan(iter func(item interface{}) bool) bool {
  578. return scan(tr.data, iter)
  579. }
  580. func scan(node *treeNode, iter func(item interface{}) bool) bool {
  581. if node.leaf {
  582. for i := 0; i < node.count; i++ {
  583. child := node.children[i]
  584. if !iter(child.unsafeItem().item) {
  585. return false
  586. }
  587. }
  588. } else {
  589. for i := 0; i < node.count; i++ {
  590. child := node.children[i]
  591. if !scan(child, iter) {
  592. return false
  593. }
  594. }
  595. }
  596. return true
  597. }
  598. // Bounds returns the bounding box of the entire R-tree
  599. func (tr *RTree) Bounds() (min, max []float64) {
  600. if tr.data.count > 0 {
  601. return tr.data.min, tr.data.max
  602. }
  603. return make([]float64, tr.dims), make([]float64, tr.dims)
  604. }
  605. // Complexity returns the complexity of the R-tree. The higher the value, the
  606. // more complex the tree. The value of 1 is the lowest.
  607. func (tr *RTree) Complexity() float64 {
  608. var nodeCount int
  609. var itemCount int
  610. tr.Traverse(func(_, _ []float64, level int, _ interface{}) bool {
  611. if level == 0 {
  612. itemCount++
  613. } else {
  614. nodeCount++
  615. }
  616. return true
  617. })
  618. return float64(tr.maxEntries*nodeCount) / float64(itemCount)
  619. }