// Package rtree - A 2d Implementation of RTree, a bounding rectangle tree. // // This file is derived from the work done by Toni Gutman. R-Trees: A Dynamic Index Structure for // Spatial Searching, Proc. 1984 ACM SIGMOD International Conference on Management of Data, pp. // 47-57. The implementation found in SQLite is a refinement of Guttman's original idea, commonly // called "R*Trees", that was described by Norbert Beckmann, Hans-Peter Kriegel, Ralf Schneider, // Bernhard Seeger: The R*-Tree: An Efficient and Robust Access Method for Points and Rectangles. // SIGMOD Conference 1990: 322-331 // // The original C code can be found at "http://www.superliminal.com/sources/sources.htm". // // And the website carries this message: "Here are a few useful bits of free source code. You're // completely free to use them for any purpose whatsoever. All I ask is that if you find one to // be particularly valuable, then consider sending feedback. Please send bugs and suggestions too. // Enjoy" package rtree import "math" // Item is an rtree item type Item interface { Rect() (minX, minY, maxX, maxY float64) } // Rect is a rectangle type Rect struct { MinX, MinY, MaxX, MaxY float64 } // Rect returns the rectangle func (item *Rect) Rect() (minX, minY, maxX, maxY float64) { return item.MinX, item.MinY, item.MaxX, item.MaxY } func min(a, b float64) float64 { if a < b { return a } return b } func max(a, b float64) float64 { if a > b { return a } return b } const ( unitSphereVolume1 = 2.000000 unitSphereVolume2 = 3.141593 unitSphereVolume3 = 4.188790 unitSphereVolume4 = 4.934802 ) const ( maxNodes = 16 minNodes = maxNodes / 2 useSphericalVolume = true unitSphereVolume = unitSphereVolume2 ) /// Minimal bounding rectangle (n-dimensional) type rectT struct { min [2]float64 ///< Min dimensions of bounding box max [2]float64 ///< Max dimensions of bounding box } /// May be data or may be another subtree /// The parents level determines this. /// If the parents level is 0, then this is data type branchT struct { rect rectT ///< Bounds child *nodeT ///< Child node item Item ///< Data ID or Ptr } /// nodeT for each branch level type nodeT struct { count int ///< Count level int ///< Leaf is zero, others positive branch [maxNodes]branchT ///< branchT } func (node *nodeT) isInternalNode() bool { return node.level > 0 } // Not a leaf, but a internal node /// A link list of nodes for reinsertion after a delete operation type listNodeT struct { next *listNodeT ///< Next in list node *nodeT ///< nodeT } /// Variables for finding a split partition type partitionVarsT struct { partition [maxNodes + 1]int total int minFill int taken [maxNodes + 1]bool count [2]int cover [2]rectT area [2]float64 branchBuf [maxNodes + 1]branchT branchCount int coverSplit rectT coverSplitArea float64 } // RTree is an implementation of an rtree type RTree struct { root *nodeT } func itemRect(item Item) (rect rectT) { minX, minY, maxX, maxY := item.Rect() return rectT{ min: [2]float64{minX, minY}, max: [2]float64{maxX, maxY}, } } // New creates a new RTree func New() *RTree { return &RTree{} } // Insert inserts item into rtree func (tr *RTree) Insert(item Item) { if tr.root == nil { tr.root = &nodeT{} } insertRect(itemRect(item), item, &tr.root, 0) } // Remove removes item from rtree func (tr *RTree) Remove(item Item) { if tr.root == nil { tr.root = &nodeT{} } removeRect(itemRect(item), item, &tr.root) } // Search finds all items in bounding box. func (tr *RTree) Search(minX, minY, maxX, maxY float64, iterator func(item Item) bool) { if iterator == nil { return } rect := rectT{ min: [2]float64{minX, minY}, max: [2]float64{maxX, maxY}, } // NOTE: May want to return search result another way, perhaps returning the number of found elements here. if tr.root == nil { tr.root = &nodeT{} } search(tr.root, rect, iterator) } // Count return the number of items in rtree. func (tr *RTree) Count() int { return countRec(tr.root, 0) } // RemoveAll removes all items from rtree. func (tr *RTree) RemoveAll() { tr.root = nil } func countRec(node *nodeT, counter int) int { if node.isInternalNode() { // not a leaf node for index := 0; index < node.count; index++ { counter = countRec(node.branch[index].child, counter) } } else { // A leaf node if node.count > 256 { println(node.count) } counter += node.count } return counter } // Inserts a new data rectangle into the index structure. // Recursively descends tree, propagates splits back up. // Returns 0 if node was not split. Old node updated. // If node was split, returns 1 and sets the pointer pointed to by // new_node to point to the new node. Old node updated to become one of two. // The level argument specifies the number of steps up from the leaf // level to insert; e.g. a data rectangle goes in at level = 0. func insertRectRec(rect rectT, item Item, node *nodeT, newNode **nodeT, level int) bool { var index int var branch branchT var otherNode *nodeT // Still above level for insertion, go down tree recursively if node.level > level { index = pickBranch(rect, node) if !insertRectRec(rect, item, node.branch[index].child, &otherNode, level) { // Child was not split node.branch[index].rect = combineRect(rect, node.branch[index].rect) return false } // Child was split node.branch[index].rect = nodeCover(node.branch[index].child) branch.child = otherNode branch.rect = nodeCover(otherNode) return addBranch(&branch, node, newNode) } else if node.level == level { // Have reached level for insertion. Add rect, split if necessary branch.rect = rect branch.item = item // Child field of leaves contains id of data record return addBranch(&branch, node, newNode) } else { // Should never occur return false } } // Insert a data rectangle into an index structure. // InsertRect provides for splitting the root; // returns 1 if root was split, 0 if it was not. // The level argument specifies the number of steps up from the leaf // level to insert; e.g. a data rectangle goes in at level = 0. // InsertRect2 does the recursion. func insertRect(rect rectT, item Item, root **nodeT, level int) bool { var newRoot *nodeT var newNode *nodeT var branch branchT if insertRectRec(rect, item, *root, &newNode, level) { // Root split newRoot = &nodeT{} // Grow tree taller and new root newRoot.level = (*root).level + 1 branch.rect = nodeCover(*root) branch.child = *root addBranch(&branch, newRoot, nil) branch.rect = nodeCover(newNode) branch.child = newNode addBranch(&branch, newRoot, nil) *root = newRoot return true } return false } // Find the smallest rectangle that includes all rectangles in branches of a node. func nodeCover(node *nodeT) rectT { var firstTime = true var rect rectT for index := 0; index < node.count; index++ { if firstTime { rect = node.branch[index].rect firstTime = false } else { rect = combineRect(rect, node.branch[index].rect) } } return rect } // Add a branch to a node. Split the node if necessary. // Returns 0 if node not split. Old node updated. // Returns 1 if node split, sets *new_node to address of new node. // Old node updated, becomes one of two. func addBranch(branch *branchT, node *nodeT, newNode **nodeT) bool { if node.count < maxNodes { // Split won't be necessary node.branch[node.count] = *branch node.count++ return false } splitNode(node, branch, newNode) return true } // Disconnect a dependent node. // Caller must return (or stop using iteration index) after this as count has changed func disconnectBranch(node *nodeT, index int) { // Remove element by swapping with the last element to prevent gaps in array node.branch[index] = node.branch[node.count-1] node.count-- } // Pick a branch. Pick the one that will need the smallest increase // in area to accommodate the new rectangle. This will result in the // least total area for the covering rectangles in the current node. // In case of a tie, pick the one which was smaller before, to get // the best resolution when searching. func pickBranch(rect rectT, node *nodeT) int { var firstTime = true var increase float64 var bestIncr float64 = -1 var area float64 var bestArea float64 var best int var tempRect rectT for index := 0; index < node.count; index++ { curRect := node.branch[index].rect area = calcRectVolume(curRect) tempRect = combineRect(rect, curRect) increase = calcRectVolume(tempRect) - area if (increase < bestIncr) || firstTime { best = index bestArea = area bestIncr = increase firstTime = false } else if (increase == bestIncr) && (area < bestArea) { best = index bestArea = area bestIncr = increase } } return best } // Combine two rectangles into larger one containing both func combineRect(rectA, rectB rectT) rectT { var newRect rectT for index := 0; index < 2; index++ { newRect.min[index] = min(rectA.min[index], rectB.min[index]) newRect.max[index] = max(rectA.max[index], rectB.max[index]) } return newRect } // Split a node. // Divides the nodes branches and the extra one between two nodes. // Old node is one of the new ones, and one really new one is created. // Tries more than one method for choosing a partition, uses best result. func splitNode(node *nodeT, branch *branchT, newNode **nodeT) { // Could just use local here, but member or external is faster since it is reused var localVars partitionVarsT var parVars = &localVars var level int // Load all the branches into a buffer, initialize old node level = node.level getBranches(node, branch, parVars) // Find partition choosePartition(parVars, minNodes) // Put branches from buffer into 2 nodes according to chosen partition *newNode = &nodeT{} node.level = level (*newNode).level = node.level loadNodes(node, *newNode, parVars) } // Calculate the n-dimensional volume of a rectangle func rectVolume(rect rectT) float64 { var volume float64 = 1 for index := 0; index < 2; index++ { volume *= rect.max[index] - rect.min[index] } return volume } // The exact volume of the bounding sphere for the given rectT func rectSphericalVolume(rect rectT) float64 { var sumOfSquares float64 var radius float64 for index := 0; index < 2; index++ { var halfExtent = (rect.max[index] - rect.min[index]) * 0.5 sumOfSquares += halfExtent * halfExtent } radius = math.Sqrt(sumOfSquares) // Pow maybe slow, so test for common dims like 2,3 and just use x*x, x*x*x. if 2 == 3 { return radius * radius * radius * unitSphereVolume } else if 2 == 2 { return radius * radius * unitSphereVolume } else { return math.Pow(radius, 2) * unitSphereVolume } } // Use one of the methods to calculate retangle volume func calcRectVolume(rect rectT) float64 { if useSphericalVolume { return rectSphericalVolume(rect) // Slower but helps certain merge cases } return rectVolume(rect) // Faster but can cause poor merges } // Load branch buffer with branches from full node plus the extra branch. func getBranches(node *nodeT, branch *branchT, parVars *partitionVarsT) { // Load the branch buffer for index := 0; index < maxNodes; index++ { parVars.branchBuf[index] = node.branch[index] } parVars.branchBuf[maxNodes] = *branch parVars.branchCount = maxNodes + 1 // Calculate rect containing all in the set parVars.coverSplit = parVars.branchBuf[0].rect for index := 1; index < maxNodes+1; index++ { parVars.coverSplit = combineRect(parVars.coverSplit, parVars.branchBuf[index].rect) } parVars.coverSplitArea = calcRectVolume(parVars.coverSplit) node.count = 0 node.level = -1 } // Method #0 for choosing a partition: // As the seeds for the two groups, pick the two rects that would waste the // most area if covered by a single rectangle, i.e. evidently the worst pair // to have in the same group. // Of the remaining, one at a time is chosen to be put in one of the two groups. // The one chosen is the one with the greatest difference in area expansion // depending on which group - the rect most strongly attracted to one group // and repelled from the other. // If one group gets too full (more would force other group to violate min // fill requirement) then other group gets the rest. // These last are the ones that can go in either group most easily. func choosePartition(parVars *partitionVarsT, minFill int) { var biggestDiff float64 var group, chosen, betterGroup int initParVars(parVars, parVars.branchCount, minFill) pickSeeds(parVars) for ((parVars.count[0] + parVars.count[1]) < parVars.total) && (parVars.count[0] < (parVars.total - parVars.minFill)) && (parVars.count[1] < (parVars.total - parVars.minFill)) { biggestDiff = -1 for index := 0; index < parVars.total; index++ { if !parVars.taken[index] { var curRect = parVars.branchBuf[index].rect rect0 := combineRect(curRect, parVars.cover[0]) rect1 := combineRect(curRect, parVars.cover[1]) growth0 := calcRectVolume(rect0) - parVars.area[0] growth1 := calcRectVolume(rect1) - parVars.area[1] diff := growth1 - growth0 if diff >= 0 { group = 0 } else { group = 1 diff = -diff } if diff > biggestDiff { biggestDiff = diff chosen = index betterGroup = group } else if (diff == biggestDiff) && (parVars.count[group] < parVars.count[betterGroup]) { chosen = index betterGroup = group } } } classify(chosen, betterGroup, parVars) } // If one group too full, put remaining rects in the other if (parVars.count[0] + parVars.count[1]) < parVars.total { if parVars.count[0] >= parVars.total-parVars.minFill { group = 1 } else { group = 0 } for index := 0; index < parVars.total; index++ { if !parVars.taken[index] { classify(index, group, parVars) } } } } // Copy branches from the buffer into two nodes according to the partition. func loadNodes(nodeA *nodeT, nodeB *nodeT, parVars *partitionVarsT) { for index := 0; index < parVars.total; index++ { if parVars.partition[index] == 0 { addBranch(&parVars.branchBuf[index], nodeA, nil) } else if parVars.partition[index] == 1 { addBranch(&parVars.branchBuf[index], nodeB, nil) } } } // Initialize a partitionVarsT structure. func initParVars(parVars *partitionVarsT, maxRects int, minFill int) { parVars.count[1] = 0 parVars.count[0] = parVars.count[1] parVars.area[1] = 0 parVars.area[0] = parVars.area[1] parVars.total = maxRects parVars.minFill = minFill for index := 0; index < maxRects; index++ { parVars.taken[index] = false parVars.partition[index] = -1 } } func pickSeeds(parVars *partitionVarsT) { var seed0, seed1 int var worst, waste float64 var area [maxNodes + 1]float64 for index := 0; index < parVars.total; index++ { area[index] = calcRectVolume(parVars.branchBuf[index].rect) } worst = -parVars.coverSplitArea - 1 for indexA := 0; indexA < parVars.total-1; indexA++ { for indexB := indexA + 1; indexB < parVars.total; indexB++ { var oneRect = combineRect(parVars.branchBuf[indexA].rect, parVars.branchBuf[indexB].rect) waste = calcRectVolume(oneRect) - area[indexA] - area[indexB] if waste > worst { worst = waste seed0 = indexA seed1 = indexB } } } classify(seed0, 0, parVars) classify(seed1, 1, parVars) } // Put a branch in one of the groups. func classify(index int, group int, parVars *partitionVarsT) { parVars.partition[index] = group parVars.taken[index] = true if parVars.count[group] == 0 { parVars.cover[group] = parVars.branchBuf[index].rect } else { parVars.cover[group] = combineRect(parVars.branchBuf[index].rect, parVars.cover[group]) } parVars.area[group] = calcRectVolume(parVars.cover[group]) parVars.count[group]++ } // Delete a data rectangle from an index structure. // Pass in a pointer to a rectT, the tid of the record, ptr to ptr to root node. // Returns 1 if record not found, 0 if success. // RemoveRect provides for eliminating the root. func removeRect(rect rectT, item Item, root **nodeT) bool { var tempNode *nodeT var reInsertList *listNodeT if !removeRectRec(rect, item, *root, &reInsertList) { // Found and deleted a data item // Reinsert any branches from eliminated nodes for reInsertList != nil { tempNode = reInsertList.node for index := 0; index < tempNode.count; index++ { insertRect(tempNode.branch[index].rect, tempNode.branch[index].item, root, tempNode.level) } reInsertList = reInsertList.next } // Check for redundant root (not leaf, 1 child) and eliminate if (*root).count == 1 && (*root).isInternalNode() { tempNode = (*root).branch[0].child *root = tempNode } return false } return true } // Delete a rectangle from non-root part of an index structure. // Called by RemoveRect. Descends tree recursively, // merges branches on the way back up. // Returns 1 if record not found, 0 if success. func removeRectRec(rect rectT, item Item, node *nodeT, listNode **listNodeT) bool { if node.isInternalNode() { // not a leaf node for index := 0; index < node.count; index++ { if overlap(rect, node.branch[index].rect) { if !removeRectRec(rect, item, node.branch[index].child, listNode) { if node.branch[index].child.count >= minNodes { // child removed, just resize parent rect node.branch[index].rect = nodeCover(node.branch[index].child) } else { // child removed, not enough entries in node, eliminate node reInsert(node.branch[index].child, listNode) disconnectBranch(node, index) // Must return after this call as count has changed } return false } } } return true } // A leaf node for index := 0; index < node.count; index++ { if node.branch[index].item == item { disconnectBranch(node, index) // Must return after this call as count has changed return false } } return true } // Decide whether two rectangles overlap. func overlap(rectA rectT, rectB rectT) bool { for index := 0; index < 2; index++ { if rectA.min[index] > rectB.max[index] || rectB.min[index] > rectA.max[index] { return false } } return true } // Add a node to the reinsertion list. All its branches will later // be reinserted into the index structure. func reInsert(node *nodeT, listNode **listNodeT) { *listNode = &listNodeT{ node: node, next: *listNode, } } // Search in an index tree or subtree for all data retangles that overlap the argument rectangle. func search(node *nodeT, rect rectT, iterator func(item Item) bool) bool { if node.isInternalNode() { // This is an internal node in the tree for index := 0; index < node.count; index++ { nrect := node.branch[index].rect if overlap(rect, nrect) { if !search(node.branch[index].child, rect, iterator) { return false // Don't continue searching } } } } else { // This is a leaf node for index := 0; index < node.count; index++ { if overlap(rect, node.branch[index].rect) { // NOTE: There are different ways to return results. Here's where to modify if !iterator(node.branch[index].item) { return false // Don't continue searching } } } } return true // Continue searching }