package base import ( "math" "unsafe" ) // precalculate infinity var mathInfNeg = math.Inf(-1) var mathInfPos = math.Inf(+1) type treeNode struct { min, max []float64 children []*treeNode count int height int leaf bool } func (node *treeNode) unsafeItem() *treeItem { return (*treeItem)(unsafe.Pointer(node)) } func (tr *RTree) createNode(children []*treeNode) *treeNode { n := &treeNode{ height: 1, leaf: true, children: make([]*treeNode, tr.maxEntries+1), } if len(children) > 0 { n.count = len(children) copy(n.children[:n.count], children) } n.min = make([]float64, tr.dims) n.max = make([]float64, tr.dims) for i := 0; i < tr.dims; i++ { n.min[i] = mathInfPos n.max[i] = mathInfNeg } return n } func (node *treeNode) extend(b *treeNode) { for i := 0; i < len(node.min); i++ { if b.min[i] < node.min[i] { node.min[i] = b.min[i] } if b.max[i] > node.max[i] { node.max[i] = b.max[i] } } } func (node *treeNode) area() float64 { area := node.max[0] - node.min[0] for i := 1; i < len(node.min); i++ { area *= node.max[i] - node.min[i] } return area } func (node *treeNode) enlargedAreaAxis(b *treeNode, axis int) float64 { var max, min float64 if b.max[axis] > node.max[axis] { max = b.max[axis] } else { max = node.max[axis] } if b.min[axis] < node.min[axis] { min = b.min[axis] } else { min = node.min[axis] } return max - min } func (node *treeNode) enlargedArea(b *treeNode) float64 { area := node.enlargedAreaAxis(b, 0) for i := 1; i < len(node.min); i++ { area *= node.enlargedAreaAxis(b, i) } return area } func (node *treeNode) intersectionAreaAxis(b *treeNode, axis int) float64 { var max, min float64 if node.max[axis] < b.max[axis] { max = node.max[axis] } else { max = b.max[axis] } if node.min[axis] > b.min[axis] { min = node.min[axis] } else { min = b.min[axis] } if max > min { return max - min } return 0 } func (node *treeNode) intersectionArea(b *treeNode) float64 { area := node.intersectionAreaAxis(b, 0) for i := 1; i < len(node.min); i++ { area *= node.intersectionAreaAxis(b, i) } return area } func (node *treeNode) margin() float64 { margin := node.max[0] - node.min[0] for i := 1; i < len(node.min); i++ { margin += node.max[i] - node.min[i] } return margin } type result int const ( not result = 0 intersects result = 1 contains result = 2 ) func (node *treeNode) overlaps(b *treeNode) result { for i := 0; i < len(node.min); i++ { if b.min[i] > node.max[i] || b.max[i] < node.min[i] { return not } if node.min[i] > b.min[i] || b.max[i] > node.max[i] { i++ for ; i < len(node.min); i++ { if b.min[i] > node.max[i] || b.max[i] < node.min[i] { return not } } return intersects } } return contains } func (node *treeNode) intersects(b *treeNode) bool { for i := 0; i < len(node.min); i++ { if b.min[i] > node.max[i] || b.max[i] < node.min[i] { return false } } return true } func (node *treeNode) findItem(item interface{}) int { for i := 0; i < node.count; i++ { if node.children[i].unsafeItem().item == item { return i } } return -1 } func (node *treeNode) contains(b *treeNode) bool { for i := 0; i < len(node.min); i++ { if node.min[i] > b.min[i] || b.max[i] > node.max[i] { return false } } return true } func (node *treeNode) childCount() int { if node.leaf { return node.count } var n int for i := 0; i < node.count; i++ { n += node.children[i].childCount() } return n } type treeItem struct { min, max []float64 item interface{} } func (item *treeItem) unsafeNode() *treeNode { return (*treeNode)(unsafe.Pointer(item)) } // RTree is an R-tree type RTree struct { dims int maxEntries int minEntries int data *treeNode // root node // resusable fields, these help performance of common mutable operations. reuse struct { path []*treeNode // for reinsertion path indexes []int // for remove function stack []int // for bulk loading } } // New creates a new R-tree func New(dims, maxEntries int) *RTree { if dims <= 0 { panic("invalid dimensions") } tr := &RTree{} tr.dims = dims tr.maxEntries = int(math.Max(4, float64(maxEntries))) tr.minEntries = int(math.Max(2, math.Ceil(float64(tr.maxEntries)*0.4))) tr.data = tr.createNode(nil) return tr } // Insert inserts an item func (tr *RTree) Insert(min, max []float64, item interface{}) { if len(min) != tr.dims || len(max) != tr.dims { panic("invalid dimensions") } if item == nil { panic("nil item") } bbox := treeNode{min: min, max: max} tr.insert(&bbox, item, tr.data.height-1, false) } func (tr *RTree) insert(bbox *treeNode, item interface{}, level int, isNode bool) { tr.reuse.path = tr.reuse.path[:0] node, insertPath := tr.chooseSubtree(bbox, tr.data, level, tr.reuse.path) if item == nil { // item is only nil when bulk loading a node if node.leaf { panic("loading node into leaf") } node.children[node.count] = bbox node.count++ } else { ti := &treeItem{min: bbox.min, max: bbox.max, item: item} node.children[node.count] = ti.unsafeNode() node.count++ } node.extend(bbox) for level >= 0 { if insertPath[level].count > tr.maxEntries { insertPath = tr.split(insertPath, level) level-- } else { break } } tr.adjustParentBBoxes(bbox, insertPath, level) tr.reuse.path = insertPath } func (tr *RTree) adjustParentBBoxes(bbox *treeNode, path []*treeNode, level int) { // adjust bboxes along the given tree path for i := level; i >= 0; i-- { path[i].extend(bbox) } } func (tr *RTree) chooseSubtree(bbox, node *treeNode, level int, path []*treeNode) (*treeNode, []*treeNode) { var targetNode *treeNode var area, enlargement, minArea, minEnlargement float64 for { path = append(path, node) if node.leaf || len(path)-1 == level { break } minEnlargement = mathInfPos minArea = minEnlargement for i := 0; i < node.count; i++ { child := node.children[i] area = child.area() enlargement = bbox.enlargedArea(child) - area if enlargement < minEnlargement { minEnlargement = enlargement if area < minArea { minArea = area } targetNode = child } else if enlargement == minEnlargement { if area < minArea { minArea = area targetNode = child } } } if targetNode != nil { node = targetNode } else if node.count > 0 { node = (*treeNode)(node.children[0]) } else { node = nil } } return node, path } func (tr *RTree) split(insertPath []*treeNode, level int) []*treeNode { var node = insertPath[level] var M = node.count var m = tr.minEntries tr.chooseSplitAxis(node, m, M) splitIndex := tr.chooseSplitIndex(node, m, M) spliced := make([]*treeNode, node.count-splitIndex) copy(spliced, node.children[splitIndex:]) node.count = splitIndex newNode := tr.createNode(spliced) newNode.height = node.height newNode.leaf = node.leaf tr.calcBBox(node) tr.calcBBox(newNode) if level != 0 { insertPath[level-1].children[insertPath[level-1].count] = newNode insertPath[level-1].count++ } else { tr.splitRoot(node, newNode) } return insertPath } func (tr *RTree) chooseSplitIndex(node *treeNode, m, M int) int { var i int var bbox1, bbox2 *treeNode var overlap, area, minOverlap, minArea float64 var index int minArea = mathInfPos minOverlap = minArea for i = m; i <= M-m; i++ { bbox1 = tr.distBBox(node, 0, i, nil) bbox2 = tr.distBBox(node, i, M, nil) overlap = bbox1.intersectionArea(bbox2) area = bbox1.area() + bbox2.area() // choose distribution with minimum overlap if overlap < minOverlap { minOverlap = overlap index = i if area < minArea { minArea = area } } else if overlap == minOverlap { // otherwise choose distribution with minimum area if area < minArea { minArea = area index = i } } } return index } func (tr *RTree) calcBBox(node *treeNode) { tr.distBBox(node, 0, node.count, node) } func (tr *RTree) chooseSplitAxis(node *treeNode, m, M int) { minMargin := tr.allDistMargin(node, m, M, 0) var minAxis int for axis := 1; axis < tr.dims; axis++ { margin := tr.allDistMargin(node, m, M, axis) if margin < minMargin { minMargin = margin minAxis = axis } } if minAxis < tr.dims { tr.sortNodes(node, minAxis) } } func (tr *RTree) splitRoot(node, newNode *treeNode) { tr.data = tr.createNode([]*treeNode{node, newNode}) tr.data.height = node.height + 1 tr.data.leaf = false tr.calcBBox(tr.data) } func (tr *RTree) distBBox(node *treeNode, k, p int, destNode *treeNode) *treeNode { if destNode == nil { destNode = tr.createNode(nil) } else { for i := 0; i < tr.dims; i++ { destNode.min[i] = mathInfPos destNode.max[i] = mathInfNeg } } for i := k; i < p; i++ { if node.leaf { destNode.extend(node.children[i]) } else { destNode.extend((*treeNode)(node.children[i])) } } return destNode } func (tr *RTree) allDistMargin(node *treeNode, m, M int, axis int) float64 { tr.sortNodes(node, axis) var leftBBox = tr.distBBox(node, 0, m, nil) var rightBBox = tr.distBBox(node, M-m, M, nil) var margin = leftBBox.margin() + rightBBox.margin() var i int if node.leaf { for i = m; i < M-m; i++ { leftBBox.extend(node.children[i]) margin += leftBBox.margin() } for i = M - m - 1; i >= m; i-- { leftBBox.extend(node.children[i]) margin += rightBBox.margin() } } else { for i = m; i < M-m; i++ { child := (*treeNode)(node.children[i]) leftBBox.extend(child) margin += leftBBox.margin() } for i = M - m - 1; i >= m; i-- { child := (*treeNode)(node.children[i]) leftBBox.extend(child) margin += rightBBox.margin() } } return margin } func (tr *RTree) sortNodes(node *treeNode, axis int) { sortByAxis(node.children[:node.count], axis) } func sortByAxis(items []*treeNode, axis int) { if len(items) < 2 { return } left, right := 0, len(items)-1 pivotIndex := len(items) / 2 items[pivotIndex], items[right] = items[right], items[pivotIndex] for i := range items { if items[i].min[axis] < items[right].min[axis] { items[i], items[left] = items[left], items[i] left++ } } items[left], items[right] = items[right], items[left] sortByAxis(items[:left], axis) sortByAxis(items[left+1:], axis) } // Search searches the tree for items in the input rectangle func (tr *RTree) Search(min, max []float64, iter func(item interface{}) bool) bool { bbox := &treeNode{min: min, max: max} if !tr.data.intersects(bbox) { return true } return tr.search(tr.data, bbox, iter) } func (tr *RTree) search(node, bbox *treeNode, iter func(item interface{}) bool) bool { if node.leaf { for i := 0; i < node.count; i++ { if bbox.intersects(node.children[i]) { if !iter(node.children[i].unsafeItem().item) { return false } } } } else { for i := 0; i < node.count; i++ { r := bbox.overlaps(node.children[i]) if r == intersects { if !tr.search(node.children[i], bbox, iter) { return false } } else if r == contains { if !scan(node.children[i], iter) { return false } } } } return true } func (tr *RTree) IsEmpty() bool { empty := true tr.Scan(func(item interface{}) bool { empty = false return false }) return empty } // Remove removes an item from the R-tree. func (tr *RTree) Remove(min, max []float64, item interface{}) { bbox := &treeNode{min: min, max: max} tr.remove(bbox, item) } func (tr *RTree) remove(bbox *treeNode, item interface{}) { path := tr.reuse.path[:0] indexes := tr.reuse.indexes[:0] var node = tr.data var i int var parent *treeNode var index int var goingUp bool for node != nil || len(path) != 0 { if node == nil { node = path[len(path)-1] path = path[:len(path)-1] if len(path) == 0 { parent = nil } else { parent = path[len(path)-1] } i = indexes[len(indexes)-1] indexes = indexes[:len(indexes)-1] goingUp = true } if node.leaf { index = node.findItem(item) if index != -1 { // item found, remove the item and condense tree upwards copy(node.children[index:], node.children[index+1:]) node.children[node.count-1] = nil node.count-- path = append(path, node) tr.condense(path) goto done } } if !goingUp && !node.leaf && node.contains(bbox) { // go down path = append(path, node) indexes = append(indexes, i) i = 0 parent = node node = (*treeNode)(node.children[0]) } else if parent != nil { // go right i++ if i == parent.count { node = nil } else { node = (*treeNode)(parent.children[i]) } goingUp = false } else { node = nil } } done: tr.reuse.path = path tr.reuse.indexes = indexes return } func (tr *RTree) condense(path []*treeNode) { // go through the path, removing empty nodes and updating bboxes var siblings []*treeNode for i := len(path) - 1; i >= 0; i-- { if path[i].count == 0 { if i > 0 { siblings = path[i-1].children[:path[i-1].count] index := -1 for j := 0; j < len(siblings); j++ { if siblings[j] == path[i] { index = j break } } copy(siblings[index:], siblings[index+1:]) siblings[len(siblings)-1] = nil path[i-1].count-- //siblings = siblings[:len(siblings)-1] //path[i-1].children = siblings } else { tr.data = tr.createNode(nil) // clear tree } } else { tr.calcBBox(path[i]) } } } // Count returns the number of items in the R-tree. func (tr *RTree) Count() int { return tr.data.childCount() } // Traverse iterates over the entire R-tree and includes all nodes and items. func (tr *RTree) Traverse(iter func(min, max []float64, level int, item interface{}) bool) bool { return tr.traverse(tr.data, iter) } func (tr *RTree) traverse(node *treeNode, iter func(min, max []float64, level int, item interface{}) bool) bool { if !iter(node.min, node.max, int(node.height), nil) { return false } if node.leaf { for i := 0; i < node.count; i++ { child := node.children[i] if !iter(child.min, child.max, 0, child.unsafeItem().item) { return false } } } else { for i := 0; i < node.count; i++ { child := node.children[i] if !tr.traverse(child, iter) { return false } } } return true } // Scan iterates over the entire R-tree func (tr *RTree) Scan(iter func(item interface{}) bool) bool { return scan(tr.data, iter) } func scan(node *treeNode, iter func(item interface{}) bool) bool { if node.leaf { for i := 0; i < node.count; i++ { child := node.children[i] if !iter(child.unsafeItem().item) { return false } } } else { for i := 0; i < node.count; i++ { child := node.children[i] if !scan(child, iter) { return false } } } return true } // Bounds returns the bounding box of the entire R-tree func (tr *RTree) Bounds() (min, max []float64) { if tr.data.count > 0 { return tr.data.min, tr.data.max } return make([]float64, tr.dims), make([]float64, tr.dims) } // Complexity returns the complexity of the R-tree. The higher the value, the // more complex the tree. The value of 1 is the lowest. func (tr *RTree) Complexity() float64 { var nodeCount int var itemCount int tr.Traverse(func(_, _ []float64, level int, _ interface{}) bool { if level == 0 { itemCount++ } else { nodeCount++ } return true }) return float64(tr.maxEntries*nodeCount) / float64(itemCount) }