当前位置: 首页 > news >正文

空间数据结构管理---RTree (下篇,代码实例)

源码请参考:GitHub - patrick-higgins/rtreego: an R-Tree library for Go

一、RTree使用说明:

#To create a new tree, specify the number of spatial dimensions and the minimum
and maximum branching factor:

    rt := rtreego.NewTree(2, 25, 50)

#You can also bulk-load the tree when creating it by passing the objects as
a parameter.

    rt := rtreego.NewTree(2, 25, 50, objects...)

#Any type that implements the `Spatial` interface can be stored in the tree:

    type Spatial interface {
      Bounds() *Rect
    }

#`Rect`s are data structures for representing spatial objects, while `Point`s
represent spatial locations.  Creating `Point`s is easy--they're just slices
of `float64`s:

    p1 := rtreego.Point{0.4, 0.5}
    p2 := rtreego.Point{6.2, -3.4}

#To create a `Rect`, specify a location and the lengths of the sides:

    r1, _ := rtreego.NewRect(p1, []float64{1, 2})
    r2, _ := rtreego.NewRect(p2, []float64{1.7, 2.7})

#To demonstrate, let's create and store some test data.

    type Thing struct {
      where *Rect
      name string
    }

    func (t *Thing) Bounds() *Rect {
      return t.where
    }

    rt.Insert(&Thing{r1, "foo"})
    rt.Insert(&Thing{r2, "bar"})

    size := rt.Size() // returns 2

#We can insert and delete objects from the tree in any order.

    rt.Delete(thing2)
    // do some stuff...
    rt.Insert(anotherThing)

#Note that ```Delete``` function does the equality comparison by comparing the
memory addresses of the objects. If you do not have a pointer to the original
object anymore, you can define a custom comparator.

    type Comparator func(obj1, obj2 Spatial) (equal bool)

#You can use a custom comparator with ```DeleteWithComparator``` function.

    cmp := func(obj1, obj2 Spatial) bool {
      sp1 := obj1.(*IDRect)
      sp2 := obj2.(*IDRect)

      return sp1.ID == sp2.ID
    }

    rt.DeleteWithComparator(obj, cmp)

#If you want to store points instead of rectangles, you can easily convert a
point into a rectangle using the `ToRect` method:

    var tol = 0.01

    type Somewhere struct {
      location rtreego.Point
      name string
      wormhole chan int
    }

    func (s *Somewhere) Bounds() *Rect {
      // define the bounds of s to be a rectangle centered at s.location
      // with side lengths 2 * tol:
      return s.location.ToRect(tol)
    }

    rt.Insert(&Somewhere{rtreego.Point{0, 0}, "Someplace", nil})

#If you want to update the location of an object, you must delete it, update it,
and re-insert.  Just modifying the object so that the `*Rect` returned by
`Location()` changes, without deleting and re-inserting the object, will
corrupt the tree.

### Queries

#Bounding-box and k-nearest-neighbors queries are supported.

#Bounding-box queries require a search `*Rect`. It returns all objects that
touch the search rectangle.

    bb, _ := rtreego.NewRect(rtreego.Point{1.7, -3.4}, []float64{3.2, 1.9})

    // Get a slice of the objects in rt that intersect bb:
    results := rt.SearchIntersect(bb)

### Filters

#You can filter out values during searches by implementing Filter functions.

    type Filter func(results []Spatial, object Spatial) (refuse, abort bool)

#A filter for limiting results by result count is included in the package for
backwards compatibility.

    // maximum of three results will be returned
    tree.SearchIntersect(bb, LimitFilter(3))

#Nearest-neighbor queries find the objects in a tree closest to a specified
query point.

    q := rtreego.Point{6.5, -2.47}
    k := 5

    // Get a slice of the k objects in rt closest to q:
    results = rt.NearestNeighbors(k, q)

### More information

#See [GoDoc](http://godoc.org/github.com/dhconnelly/rtreego) for full API
documentation.

二、源码

rtree.go

package rtreego

import (
	"fmt"
	"math"
	"sort"
)

// Comparator compares two spatials and returns whether they are equal.
type Comparator func(obj1, obj2 Spatial) (equal bool)

func defaultComparator(obj1, obj2 Spatial) bool {
	return obj1 == obj2
}

// Rtree represents an R-tree, a balanced search tree for storing and querying
// spatial objects.  Dim specifies the number of spatial dimensions and
// MinChildren/MaxChildren specify the minimum/maximum branching factors.
type Rtree struct {
	Dim         int
	MinChildren int
	MaxChildren int
	root        *node
	size        int
	height      int
}

// NewTree returns an Rtree. If the number of objects given on initialization
// is larger than max, the Rtree will be initialized using the Overlap
// Minimizing Top-down bulk-loading algorithm.
func NewTree(dim, min, max int, objs ...Spatial) *Rtree {
	rt := &Rtree{
		Dim:         dim,
		MinChildren: min,
		MaxChildren: max,
		height:      1,
		root: &node{
			entries: []entry{},
			leaf:    true,
			level:   1,
		},
	}

	if len(objs) <= rt.MaxChildren {
		for _, obj := range objs {
			rt.Insert(obj)
		}
	} else {
		rt.bulkLoad(objs)
	}

	return rt
}

// Size returns the number of objects currently stored in tree.
func (tree *Rtree) Size() int {
	return tree.size
}

func (tree *Rtree) String() string {
	return "foo"
}

// Depth returns the maximum depth of tree.
func (tree *Rtree) Depth() int {
	return tree.height
}

type dimSorter struct {
	dim  int
	objs []entry
}

func (s *dimSorter) Len() int {
	return len(s.objs)
}

func (s *dimSorter) Swap(i, j int) {
	s.objs[i], s.objs[j] = s.objs[j], s.objs[i]
}

func (s *dimSorter) Less(i, j int) bool {
	return s.objs[i].bb.p[s.dim] < s.objs[j].bb.p[s.dim]
}

// walkPartitions splits objs into slices of maximum k elements and
// iterates over these partitions.
func walkPartitions(k int, objs []entry, iter func(parts []entry)) {
	n := (len(objs) + k - 1) / k // ceil(len(objs) / k)

	for i := 1; i < n; i++ {
		iter(objs[(i-1)*k : i*k])
	}
	iter(objs[(n-1)*k:])
}

func sortByDim(dim int, objs []entry) {
	sort.Sort(&dimSorter{dim, objs})
}

// bulkLoad bulk loads the Rtree using OMT algorithm. bulkLoad contains special
// handling for the root node.
func (tree *Rtree) bulkLoad(objs []Spatial) {
	n := len(objs)

	// create entries for all the objects
	entries := make([]entry, n)
	for i := range objs {
		entries[i] = entry{
			bb:  objs[i].Bounds(),
			obj: objs[i],
		}
	}

	// following equations are defined in the paper describing OMT
	var (
		N = float64(n)
		M = float64(tree.MaxChildren)
	)
	// Eq1: height of the tree
	// use log2 instead of log due to rounding errors with log,
	// eg, math.Log(9) / math.Log(3) > 2
	h := math.Ceil(math.Log2(N) / math.Log2(M))

	// Eq2: size of subtrees at the root
	nsub := math.Pow(M, h-1)

	// Inner Eq3: number of subtrees at the root
	s := math.Ceil(N / nsub)

	// Eq3: number of slices
	S := math.Floor(math.Sqrt(s))

	// sort all entries by first dimension
	sortByDim(0, entries)

	tree.height = int(h)
	tree.size = n
	tree.root = tree.omt(int(h), int(S), entries, int(s))
}

// omt is the recursive part of the Overlap Minimizing Top-loading bulk-
// load approach. Returns the root node of a subtree.
func (tree *Rtree) omt(level, nSlices int, objs []entry, m int) *node {
	// if number of objects is less than or equal than max children per leaf,
	// we need to create a leaf node
	if len(objs) <= m {
		// as long as the recursion is not at the leaf, call it again
		if level > 1 {
			child := tree.omt(level-1, nSlices, objs, m)
			n := &node{
				level: level,
				entries: []entry{{
					bb:    child.computeBoundingBox(),
					child: child,
				}},
			}
			child.parent = n
			return n
		}
		return &node{
			leaf:    true,
			entries: objs,
			level:   level,
		}
	}

	n := &node{
		level:   level,
		entries: make([]entry, 0, m),
	}

	// maximum node size given at most M nodes at this level
	k := (len(objs) + m - 1) / m // = ceil(N / M)

	// In the root level, split objs in nSlices. In all other levels,
	// we use a single slice.
	vertSize := len(objs)
	if nSlices > 1 {
		vertSize = nSlices * k
	}

	// create sub trees
	walkPartitions(vertSize, objs, func(vert []entry) {
		// sort vertical slice by a different dimension on every level
		sortByDim((tree.height-level+1)%tree.Dim, vert)

		// split slice into groups of size k
		walkPartitions(k, vert, func(part []entry) {
			child := tree.omt(level-1, 1, part, tree.MaxChildren)
			child.parent = n

			n.entries = append(n.entries, entry{
				bb:    child.computeBoundingBox(),
				child: child,
			})
		})
	})
	return n
}

// node represents a tree node of an Rtree.
type node struct {
	parent  *node
	leaf    bool
	entries []entry
	level   int // node depth in the Rtree
}

func (n *node) String() string {
	return fmt.Sprintf("node{leaf: %v, entries: %v}", n.leaf, n.entries)
}

// entry represents a spatial index record stored in a tree node.
type entry struct {
	bb    *Rect // bounding-box of all children of this entry
	child *node
	obj   Spatial
}

func (e entry) String() string {
	if e.child != nil {
		return fmt.Sprintf("entry{bb: %v, child: %v}", e.bb, e.child)
	}
	return fmt.Sprintf("entry{bb: %v, obj: %v}", e.bb, e.obj)
}

// Spatial is an interface for objects that can be stored in an Rtree and queried.
type Spatial interface {
	Bounds() *Rect
}

// Insertion

// Insert inserts a spatial object into the tree.  If insertion
// causes a leaf node to overflow, the tree is rebalanced automatically.
//
// Implemented per Section 3.2 of "R-trees: A Dynamic Index Structure for
// Spatial Searching" by A. Guttman, Proceedings of ACM SIGMOD, p. 47-57, 1984.
func (tree *Rtree) Insert(obj Spatial) {
	e := entry{obj.Bounds(), nil, obj}
	tree.insert(e, 1)
	tree.size++
}

// insert adds the specified entry to the tree at the specified level.
func (tree *Rtree) insert(e entry, level int) {
	leaf := tree.chooseNode(tree.root, e, level)
	leaf.entries = append(leaf.entries, e)

	// update parent pointer if necessary
	if e.child != nil {
		e.child.parent = leaf
	}

	// split leaf if overflows
	var split *node
	if len(leaf.entries) > tree.MaxChildren {
		leaf, split = leaf.split(tree.MinChildren)
	}
	root, splitRoot := tree.adjustTree(leaf, split)
	if splitRoot != nil {
		oldRoot := root
		tree.height++
		tree.root = &node{
			parent: nil,
			level:  tree.height,
			entries: []entry{
				{bb: oldRoot.computeBoundingBox(), child: oldRoot},
				{bb: splitRoot.computeBoundingBox(), child: splitRoot},
			},
		}
		oldRoot.parent = tree.root
		splitRoot.parent = tree.root
	}
}

// chooseNode finds the node at the specified level to which e should be added.
func (tree *Rtree) chooseNode(n *node, e entry, level int) *node {
	if n.leaf || n.level == level {
		return n
	}

	// find the entry whose bb needs least enlargement to include obj
	diff := math.MaxFloat64
	var chosen entry
	for _, en := range n.entries {
		bb := boundingBox(en.bb, e.bb)
		d := bb.Size() - en.bb.Size()
		if d < diff || (d == diff && en.bb.Size() < chosen.bb.Size()) {
			diff = d
			chosen = en
		}
	}

	return tree.chooseNode(chosen.child, e, level)
}

// adjustTree splits overflowing nodes and propagates the changes upwards.
func (tree *Rtree) adjustTree(n, nn *node) (*node, *node) {
	// Let the caller handle root adjustments.
	if n == tree.root {
		return n, nn
	}

	// Re-size the bounding box of n to account for lower-level changes.
	en := n.getEntry()
	en.bb = n.computeBoundingBox()

	// If nn is nil, then we're just propagating changes upwards.
	if nn == nil {
		return tree.adjustTree(n.parent, nil)
	}

	// Otherwise, these are two nodes resulting from a split.
	// n was reused as the "left" node, but we need to add nn to n.parent.
	enn := entry{nn.computeBoundingBox(), nn, nil}
	n.parent.entries = append(n.parent.entries, enn)

	// If the new entry overflows the parent, split the parent and propagate.
	if len(n.parent.entries) > tree.MaxChildren {
		return tree.adjustTree(n.parent.split(tree.MinChildren))
	}

	// Otherwise keep propagating changes upwards.
	return tree.adjustTree(n.parent, nil)
}

// getEntry returns a pointer to the entry for the node n from n's parent.
func (n *node) getEntry() *entry {
	var e *entry
	for i := range n.parent.entries {
		if n.parent.entries[i].child == n {
			e = &n.parent.entries[i]
			break
		}
	}
	return e
}

// computeBoundingBox finds the MBR of the children of n.
func (n *node) computeBoundingBox() (bb *Rect) {
	childBoxes := make([]*Rect, len(n.entries))
	for i, e := range n.entries {
		childBoxes[i] = e.bb
	}
	bb = boundingBoxN(childBoxes...)
	return
}

// split splits a node into two groups while attempting to minimize the
// bounding-box area of the resulting groups.
func (n *node) split(minGroupSize int) (left, right *node) {
	// find the initial split
	l, r := n.pickSeeds()
	leftSeed, rightSeed := n.entries[l], n.entries[r]

	// get the entries to be divided between left and right
	remaining := append(n.entries[:l], n.entries[l+1:r]...)
	remaining = append(remaining, n.entries[r+1:]...)

	// setup the new split nodes, but re-use n as the left node
	left = n
	left.entries = []entry{leftSeed}
	right = &node{
		parent:  n.parent,
		leaf:    n.leaf,
		level:   n.level,
		entries: []entry{rightSeed},
	}

	// TODO
	if rightSeed.child != nil {
		rightSeed.child.parent = right
	}
	if leftSeed.child != nil {
		leftSeed.child.parent = left
	}

	// distribute all of n's old entries into left and right.
	for len(remaining) > 0 {
		next := pickNext(left, right, remaining)
		e := remaining[next]

		if len(remaining)+len(left.entries) <= minGroupSize {
			assign(e, left)
		} else if len(remaining)+len(right.entries) <= minGroupSize {
			assign(e, right)
		} else {
			assignGroup(e, left, right)
		}

		remaining = append(remaining[:next], remaining[next+1:]...)
	}

	return
}

// getAllBoundingBoxes traverses tree populating slice of bounding boxes of non-leaf nodes.
func (n *node) getAllBoundingBoxes() []*Rect {
	var rects []*Rect
	if n.leaf {
		return rects
	}
	for _, e := range n.entries {
		if e.child == nil {
			return rects
		}
		rectsInter := append(e.child.getAllBoundingBoxes(), e.bb)
		rects = append(rects, rectsInter...)
	}
	return rects
}

func assign(e entry, group *node) {
	if e.child != nil {
		e.child.parent = group
	}
	group.entries = append(group.entries, e)
}

// assignGroup chooses one of two groups to which a node should be added.
func assignGroup(e entry, left, right *node) {
	leftBB := left.computeBoundingBox()
	rightBB := right.computeBoundingBox()
	leftEnlarged := boundingBox(leftBB, e.bb)
	rightEnlarged := boundingBox(rightBB, e.bb)

	// first, choose the group that needs the least enlargement
	leftDiff := leftEnlarged.Size() - leftBB.Size()
	rightDiff := rightEnlarged.Size() - rightBB.Size()
	if diff := leftDiff - rightDiff; diff < 0 {
		assign(e, left)
		return
	} else if diff > 0 {
		assign(e, right)
		return
	}

	// next, choose the group that has smaller area
	if diff := leftBB.Size() - rightBB.Size(); diff < 0 {
		assign(e, left)
		return
	} else if diff > 0 {
		assign(e, right)
		return
	}

	// next, choose the group with fewer entries
	if diff := len(left.entries) - len(right.entries); diff <= 0 {
		assign(e, left)
		return
	}
	assign(e, right)
}

// pickSeeds chooses two child entries of n to start a split.
func (n *node) pickSeeds() (int, int) {
	left, right := 0, 1
	maxWastedSpace := -1.0
	for i, e1 := range n.entries {
		for j, e2 := range n.entries[i+1:] {
			d := boundingBox(e1.bb, e2.bb).Size() - e1.bb.Size() - e2.bb.Size()
			if d > maxWastedSpace {
				maxWastedSpace = d
				left, right = i, j+i+1
			}
		}
	}
	return left, right
}

// pickNext chooses an entry to be added to an entry group.
func pickNext(left, right *node, entries []entry) (next int) {
	maxDiff := -1.0
	leftBB := left.computeBoundingBox()
	rightBB := right.computeBoundingBox()
	for i, e := range entries {
		d1 := boundingBox(leftBB, e.bb).Size() - leftBB.Size()
		d2 := boundingBox(rightBB, e.bb).Size() - rightBB.Size()
		d := math.Abs(d1 - d2)
		if d > maxDiff {
			maxDiff = d
			next = i
		}
	}
	return
}

// Deletion

// Delete removes an object from the tree.  If the object is not found, returns
// false, otherwise returns true. Uses the default comparator when checking
// equality.
//
// Implemented per Section 3.3 of "R-trees: A Dynamic Index Structure for
// Spatial Searching" by A. Guttman, Proceedings of ACM SIGMOD, p. 47-57, 1984.
func (tree *Rtree) Delete(obj Spatial) bool {
	return tree.DeleteWithComparator(obj, defaultComparator)
}

// DeleteWithComparator removes an object from the tree using a custom
// comparator for evaluating equalness. This is useful when you want to remove
// an object from a tree but don't have a pointer to the original object
// anymore.
func (tree *Rtree) DeleteWithComparator(obj Spatial, cmp Comparator) bool {
	n := tree.findLeaf(tree.root, obj, cmp)
	if n == nil {
		return false
	}

	ind := -1
	for i, e := range n.entries {
		if cmp(e.obj, obj) {
			ind = i
		}
	}
	if ind < 0 {
		return false
	}

	n.entries = append(n.entries[:ind], n.entries[ind+1:]...)

	tree.condenseTree(n)
	tree.size--

	if !tree.root.leaf && len(tree.root.entries) == 1 {
		tree.root = tree.root.entries[0].child
	}

	tree.height = tree.root.level

	return true
}

// findLeaf finds the leaf node containing obj.
func (tree *Rtree) findLeaf(n *node, obj Spatial, cmp Comparator) *node {
	if n.leaf {
		return n
	}
	// if not leaf, search all candidate subtrees
	for _, e := range n.entries {
		if e.bb.containsRect(obj.Bounds()) {
			leaf := tree.findLeaf(e.child, obj, cmp)
			if leaf == nil {
				continue
			}
			// check if the leaf actually contains the object
			for _, leafEntry := range leaf.entries {
				if cmp(leafEntry.obj, obj) {
					return leaf
				}
			}
		}
	}
	return nil
}

// condenseTree deletes underflowing nodes and propagates the changes upwards.
func (tree *Rtree) condenseTree(n *node) {
	deleted := []*node{}

	for n != tree.root {
		if len(n.entries) < tree.MinChildren {
			// remove n from parent entries
			entries := []entry{}
			for _, e := range n.parent.entries {
				if e.child != n {
					entries = append(entries, e)
				}
			}
			if len(n.parent.entries) == len(entries) {
				panic(fmt.Errorf("Failed to remove entry from parent"))
			}
			n.parent.entries = entries

			// only add n to deleted if it still has children
			if len(n.entries) > 0 {
				deleted = append(deleted, n)
			}
		} else {
			// just a child entry deletion, no underflow
			n.getEntry().bb = n.computeBoundingBox()
		}
		n = n.parent
	}

	for _, n := range deleted {
		// reinsert entry so that it will remain at the same level as before
		e := entry{n.computeBoundingBox(), n, nil}
		tree.insert(e, n.level+1)
	}
}

// Searching

// SearchIntersect returns all objects that intersect the specified rectangle.
// Implemented per Section 3.1 of "R-trees: A Dynamic Index Structure for
// Spatial Searching" by A. Guttman, Proceedings of ACM SIGMOD, p. 47-57, 1984.
func (tree *Rtree) SearchIntersect(bb *Rect, filters ...Filter) []Spatial {
	return tree.searchIntersect([]Spatial{}, tree.root, bb, filters)
}

// SearchIntersectWithLimit is similar to SearchIntersect, but returns
// immediately when the first k results are found. A negative k behaves exactly
// like SearchIntersect and returns all the results.
//
// Kept for backwards compatibility, please use SearchIntersect with a
// LimitFilter.
func (tree *Rtree) SearchIntersectWithLimit(k int, bb *Rect) []Spatial {
	// backwards compatibility, previous implementation didn't limit results if
	// k was negative.
	if k < 0 {
		return tree.SearchIntersect(bb)
	}
	return tree.SearchIntersect(bb, LimitFilter(k))
}

func (tree *Rtree) searchIntersect(results []Spatial, n *node, bb *Rect, filters []Filter) []Spatial {
	for _, e := range n.entries {
		if !intersect(e.bb, bb) {
			continue
		}

		if !n.leaf {
			results = tree.searchIntersect(results, e.child, bb, filters)
			continue
		}

		refuse, abort := applyFilters(results, e.obj, filters)
		if !refuse {
			results = append(results, e.obj)
		}

		if abort {
			break
		}
	}
	return results
}

// NearestNeighbor returns the closest object to the specified point.
// Implemented per "Nearest Neighbor Queries" by Roussopoulos et al
func (tree *Rtree) NearestNeighbor(p Point) Spatial {
	obj, _ := tree.nearestNeighbor(p, tree.root, math.MaxFloat64, nil)
	return obj
}

// GetAllBoundingBoxes returning slice of bounding boxes by traversing tree. Slice
// includes bounding boxes from all non-leaf nodes.
func (tree *Rtree) GetAllBoundingBoxes() []*Rect {
	var rects []*Rect
	if tree.root != nil {
		rects = tree.root.getAllBoundingBoxes()
	}
	return rects
}

// utilities for sorting slices of entries

type entrySlice struct {
	entries []entry
	dists   []float64
}

func (s entrySlice) Len() int { return len(s.entries) }

func (s entrySlice) Swap(i, j int) {
	s.entries[i], s.entries[j] = s.entries[j], s.entries[i]
	s.dists[i], s.dists[j] = s.dists[j], s.dists[i]
}

func (s entrySlice) Less(i, j int) bool {
	return s.dists[i] < s.dists[j]
}

func sortEntries(p Point, entries []entry) ([]entry, []float64) {
	sorted := make([]entry, len(entries))
	dists := make([]float64, len(entries))
	return sortPreallocEntries(p, entries, sorted, dists)
}

func sortPreallocEntries(p Point, entries, sorted []entry, dists []float64) ([]entry, []float64) {
	// use preallocated slices
	sorted = sorted[:len(entries)]
	dists = dists[:len(entries)]

	for i := 0; i < len(entries); i++ {
		sorted[i] = entries[i]
		dists[i] = p.minDist(entries[i].bb)
	}
	sort.Sort(entrySlice{sorted, dists})
	return sorted, dists
}

func pruneEntries(p Point, entries []entry, minDists []float64) []entry {
	minMinMaxDist := math.MaxFloat64
	for i := range entries {
		minMaxDist := p.minMaxDist(entries[i].bb)
		if minMaxDist < minMinMaxDist {
			minMinMaxDist = minMaxDist
		}
	}
	// remove all entries with minDist > minMinMaxDist
	pruned := []entry{}
	for i := range entries {
		if minDists[i] <= minMinMaxDist {
			pruned = append(pruned, entries[i])
		}
	}
	return pruned
}

func pruneEntriesMinDist(d float64, entries []entry, minDists []float64) []entry {
	var i int
	for ; i < len(entries); i++ {
		if minDists[i] > d {
			break
		}
	}
	return entries[:i]
}

func (tree *Rtree) nearestNeighbor(p Point, n *node, d float64, nearest Spatial) (Spatial, float64) {
	if n.leaf {
		for _, e := range n.entries {
			dist := math.Sqrt(p.minDist(e.bb))
			if dist < d {
				d = dist
				nearest = e.obj
			}
		}
	} else {
		branches, dists := sortEntries(p, n.entries)
		branches = pruneEntries(p, branches, dists)
		for _, e := range branches {
			subNearest, dist := tree.nearestNeighbor(p, e.child, d, nearest)
			if dist < d {
				d = dist
				nearest = subNearest
			}
		}
	}

	return nearest, d
}

// NearestNeighbors gets the closest Spatials to the Point.
func (tree *Rtree) NearestNeighbors(k int, p Point, filters ...Filter) []Spatial {
	// preallocate the buffers for sortings the branches. At each level of the
	// tree, we slide the buffer by the number of entries in the node.
	maxBufSize := tree.MaxChildren * tree.Depth()
	branches := make([]entry, maxBufSize)
	branchDists := make([]float64, maxBufSize)

	// allocate the buffers for the results
	dists := make([]float64, 0, k)
	objs := make([]Spatial, 0, k)

	objs, _, _ = tree.nearestNeighbors(k, p, tree.root, dists, objs, filters, branches, branchDists)
	return objs
}

// insert obj into nearest and return the first k elements in increasing order.
func insertNearest(k int, dists []float64, nearest []Spatial, dist float64, obj Spatial, filters []Filter) ([]float64, []Spatial, bool) {
	i := sort.SearchFloat64s(dists, dist)
	for i < len(nearest) && dist >= dists[i] {
		i++
	}
	if i >= k {
		return dists, nearest, false
	}

	if refuse, abort := applyFilters(nearest, obj, filters); refuse || abort {
		return dists, nearest, abort
	}

	// no resize since cap = k
	if len(nearest) < k {
		dists = append(dists, 0)
		nearest = append(nearest, nil)
	}

	left, right := dists[:i], dists[i:len(dists)-1]
	copy(dists, left)
	copy(dists[i+1:], right)
	dists[i] = dist

	leftObjs, rightObjs := nearest[:i], nearest[i:len(nearest)-1]
	copy(nearest, leftObjs)
	copy(nearest[i+1:], rightObjs)
	nearest[i] = obj

	return dists, nearest, false
}

func (tree *Rtree) nearestNeighbors(k int, p Point, n *node, dists []float64, nearest []Spatial, filters []Filter, b []entry, bd []float64) ([]Spatial, []float64, bool) {
	var abort bool
	if n.leaf {
		for _, e := range n.entries {
			dist := p.minDist(e.bb)
			dists, nearest, abort = insertNearest(k, dists, nearest, dist, e.obj, filters)
			if abort {
				break
			}
		}
	} else {
		branches, branchDists := sortPreallocEntries(p, n.entries, b, bd)
		// only prune if buffer has k elements
		if l := len(dists); l >= k {
			branches = pruneEntriesMinDist(dists[l-1], branches, branchDists)
		}
		for _, e := range branches {
			nearest, dists, abort = tree.nearestNeighbors(k, p, e.child, dists, nearest, filters, b[len(n.entries):], bd[len(n.entries):])
			if abort {
				break
			}
		}
	}
	return nearest, dists, abort
}

geom.go

package rtreego

import (
	"fmt"
	"math"
	"strings"
)

// DimError represents a failure due to mismatched dimensions.
type DimError struct {
	Expected int
	Actual   int
}

func (err DimError) Error() string {
	return "rtreego: dimension mismatch"
}

// DistError is an improper distance measurement.  It implements the error
// and is generated when a distance-related assertion fails.
type DistError float64

func (err DistError) Error() string {
	return "rtreego: improper distance"
}

// Point represents a point in n-dimensional Euclidean space.
type Point []float64

// Dist computes the Euclidean distance between two points p and q.
func (p Point) dist(q Point) float64 {
	if len(p) != len(q) {
		panic(DimError{len(p), len(q)})
	}
	sum := 0.0
	for i := range p {
		dx := p[i] - q[i]
		sum += dx * dx
	}
	return math.Sqrt(sum)
}

// minDist computes the square of the distance from a point to a rectangle.
// If the point is contained in the rectangle then the distance is zero.
//
// Implemented per Definition 2 of "Nearest Neighbor Queries" by
// N. Roussopoulos, S. Kelley and F. Vincent, ACM SIGMOD, pages 71-79, 1995.
func (p Point) minDist(r *Rect) float64 {
	if len(p) != len(r.p) {
		panic(DimError{len(p), len(r.p)})
	}

	sum := 0.0
	for i, pi := range p {
		if pi < r.p[i] {
			d := pi - r.p[i]
			sum += d * d
		} else if pi > r.q[i] {
			d := pi - r.q[i]
			sum += d * d
		} else {
			sum += 0
		}
	}
	return sum
}

// minMaxDist computes the minimum of the maximum distances from p to points
// on r.  If r is the bounding box of some geometric objects, then there is
// at least one object contained in r within minMaxDist(p, r) of p.
//
// Implemented per Definition 4 of "Nearest Neighbor Queries" by
// N. Roussopoulos, S. Kelley and F. Vincent, ACM SIGMOD, pages 71-79, 1995.
func (p Point) minMaxDist(r *Rect) float64 {
	if len(p) != len(r.p) {
		panic(DimError{len(p), len(r.p)})
	}

	// by definition, MinMaxDist(p, r) =
	// min{1<=k<=n}(|pk - rmk|^2 + sum{1<=i<=n, i != k}(|pi - rMi|^2))
	// where rmk and rMk are defined as follows:

	rm := func(k int) float64 {
		if p[k] <= (r.p[k]+r.q[k])/2 {
			return r.p[k]
		}
		return r.q[k]
	}

	rM := func(k int) float64 {
		if p[k] >= (r.p[k]+r.q[k])/2 {
			return r.p[k]
		}
		return r.q[k]
	}

	// This formula can be computed in linear time by precomputing
	// S = sum{1<=i<=n}(|pi - rMi|^2).

	S := 0.0
	for i := range p {
		d := p[i] - rM(i)
		S += d * d
	}

	// Compute MinMaxDist using the precomputed S.
	min := math.MaxFloat64
	for k := range p {
		d1 := p[k] - rM(k)
		d2 := p[k] - rm(k)
		d := S - d1*d1 + d2*d2
		if d < min {
			min = d
		}
	}

	return min
}

// Rect represents a subset of n-dimensional Euclidean space of the form
// [a1, b1] x [a2, b2] x ... x [an, bn], where ai < bi for all 1 <= i <= n.
type Rect struct {
	p, q Point // Enforced by NewRect: p[i] <= q[i] for all i.
}

// PointCoord returns the coordinate of the point of the rectangle at i
func (r *Rect) PointCoord(i int) float64 {
	return r.p[i]
}

// LengthsCoord returns the coordinate of the lengths of the rectangle at i
func (r *Rect) LengthsCoord(i int) float64 {
	return r.q[i] - r.p[i]
}

// Equal returns true if the two rectangles are equal
func (r *Rect) Equal(other *Rect) bool {
	for i, e := range r.p {
		if e != other.p[i] {
			return false
		}
	}
	for i, e := range r.q {
		if e != other.q[i] {
			return false
		}
	}
	return true
}

func (r *Rect) String() string {
	s := make([]string, len(r.p))
	for i, a := range r.p {
		b := r.q[i]
		s[i] = fmt.Sprintf("[%.2f, %.2f]", a, b)
	}
	return strings.Join(s, "x")
}

// NewRect constructs and returns a pointer to a Rect given a corner point and
// the lengths of each dimension.  The point p should be the most-negative point
// on the rectangle (in every dimension) and every length should be positive.
func NewRect(p Point, lengths []float64) (r *Rect, err error) {
	r = new(Rect)
	r.p = p
	if len(p) != len(lengths) {
		err = &DimError{len(p), len(lengths)}
		return
	}
	r.q = make([]float64, len(p))
	for i := range p {
		if lengths[i] <= 0 {
			err = DistError(lengths[i])
			return
		}
		r.q[i] = p[i] + lengths[i]
	}
	return
}

// NewRectFromPoints constructs and returns a pointer to a Rect given a corner points.
func NewRectFromPoints(minPoint, maxPoint Point) (r *Rect, err error) {
	if len(minPoint) != len(maxPoint) {
		err = &DimError{len(minPoint), len(maxPoint)}
		return
	}

	//checking that  min and max points is swapping
	for i, p := range minPoint {
		if minPoint[i] > maxPoint[i] {
			minPoint[i] = maxPoint[i]
			maxPoint[i] = p
		}
	}

	r = &Rect{p: minPoint, q: maxPoint}
	return
}

// Size computes the measure of a rectangle (the product of its side lengths).
func (r *Rect) Size() float64 {
	size := 1.0
	for i, a := range r.p {
		b := r.q[i]
		size *= b - a
	}
	return size
}

// margin computes the sum of the edge lengths of a rectangle.
func (r *Rect) margin() float64 {
	// The number of edges in an n-dimensional rectangle is n * 2^(n-1)
	// (http://en.wikipedia.org/wiki/Hypercube_graph).  Thus the number
	// of edges of length (ai - bi), where the rectangle is determined
	// by p = (a1, a2, ..., an) and q = (b1, b2, ..., bn), is 2^(n-1).
	//
	// The margin of the rectangle, then, is given by the formula
	// 2^(n-1) * [(b1 - a1) + (b2 - a2) + ... + (bn - an)].
	dim := len(r.p)
	sum := 0.0
	for i, a := range r.p {
		b := r.q[i]
		sum += b - a
	}
	return math.Pow(2, float64(dim-1)) * sum
}

// containsPoint tests whether p is located inside or on the boundary of r.
func (r *Rect) containsPoint(p Point) bool {
	if len(p) != len(r.p) {
		panic(DimError{len(r.p), len(p)})
	}

	for i, a := range p {
		// p is contained in (or on) r if and only if p <= a <= q for
		// every dimension.
		if a < r.p[i] || a > r.q[i] {
			return false
		}
	}

	return true
}

// containsRect tests whether r2 is is located inside r1.
func (r *Rect) containsRect(r2 *Rect) bool {
	if len(r.p) != len(r2.p) {
		panic(DimError{len(r.p), len(r2.p)})
	}

	for i, a1 := range r.p {
		b1, a2, b2 := r.q[i], r2.p[i], r2.q[i]
		// enforced by constructor: a1 <= b1 and a2 <= b2.
		// so containment holds if and only if a1 <= a2 <= b2 <= b1
		// for every dimension.
		if a1 > a2 || b2 > b1 {
			return false
		}
	}

	return true
}

// intersect computes the intersection of two rectangles.  If no intersection
// exists, the intersection is nil.
func intersect(r1, r2 *Rect) bool {
	dim := len(r1.p)
	if len(r2.p) != dim {
		panic(DimError{dim, len(r2.p)})
	}

	// There are four cases of overlap:
	//
	//     1.  a1------------b1
	//              a2------------b2
	//              p--------q
	//
	//     2.       a1------------b1
	//         a2------------b2
	//              p--------q
	//
	//     3.  a1-----------------b1
	//              a2-------b2
	//              p--------q
	//
	//     4.       a1-------b1
	//         a2-----------------b2
	//              p--------q
	//
	// Thus there are only two cases of non-overlap:
	//
	//     1. a1------b1
	//                    a2------b2
	//
	//     2.             a1------b1
	//        a2------b2
	//
	// Enforced by constructor: a1 <= b1 and a2 <= b2.  So we can just
	// check the endpoints.

	for i := range r1.p {
		a1, b1, a2, b2 := r1.p[i], r1.q[i], r2.p[i], r2.q[i]
		if b2 <= a1 || b1 <= a2 {
			return false
		}
	}
	return true
}

// ToRect constructs a rectangle containing p with side lengths 2*tol.
func (p Point) ToRect(tol float64) *Rect {
	dim := len(p)
	a, b := make([]float64, dim), make([]float64, dim)
	for i := range p {
		a[i] = p[i] - tol
		b[i] = p[i] + tol
	}
	return &Rect{a, b}
}

// boundingBox constructs the smallest rectangle containing both r1 and r2.
func boundingBox(r1, r2 *Rect) (bb *Rect) {
	bb = new(Rect)
	dim := len(r1.p)
	bb.p = make([]float64, dim)
	bb.q = make([]float64, dim)
	if len(r2.p) != dim {
		panic(DimError{dim, len(r2.p)})
	}
	for i := 0; i < dim; i++ {
		if r1.p[i] <= r2.p[i] {
			bb.p[i] = r1.p[i]
		} else {
			bb.p[i] = r2.p[i]
		}
		if r1.q[i] <= r2.q[i] {
			bb.q[i] = r2.q[i]
		} else {
			bb.q[i] = r1.q[i]
		}
	}
	return
}

// boundingBoxN constructs the smallest rectangle containing all of r...
func boundingBoxN(rects ...*Rect) (bb *Rect) {
	if len(rects) == 1 {
		bb = rects[0]
		return
	}
	bb = boundingBox(rects[0], rects[1])
	for _, rect := range rects[2:] {
		bb = boundingBox(bb, rect)
	}
	return
}

filter.go

package rtreego

// Filter is an interface for filtering leaves during search. The parameters
// should be treated as read-only. If refuse is true, the current entry will
// not be added to the result set. If abort is true, the search is aborted and
// the current result set will be returned.
type Filter func(results []Spatial, object Spatial) (refuse, abort bool)

// ApplyFilters applies the given filters and returns whether the entry is
// refused and/or the search should be aborted. If a filter refuses an entry,
// the following filters are not applied for the entry. If a filter aborts, the
// search terminates without further applying any filter.
func applyFilters(results []Spatial, object Spatial, filters []Filter) (bool, bool) {
	for _, filter := range filters {
		refuse, abort := filter(results, object)
		if refuse || abort {
			return refuse, abort
		}
	}
	return false, false
}

// LimitFilter checks if the results have reached the limit size and aborts if so.
func LimitFilter(limit int) Filter {
	return func(results []Spatial, object Spatial) (refuse, abort bool) {
		if len(results) >= limit {
			return true, true
		}

		return false, false
	}
}

相关文章:

  • 看我炫一下
  • Array.from(new Set)去重 与Array.map()
  • Go 学习笔记(89) — 接口类型变量的等值比较操作(nil 接口变量、空接口类型变量、非空接口类型变量)
  • dubbo源码解析之服务调用(通信)流程
  • Linux网络技术学习(四)—— 用户空间与内核的接口
  • Django--ORM 多表查询
  • pytest 运行方式、常用参数、前后置条件
  • MySQL-1-SQL讲解
  • 数据结构与算法之美读书笔记15
  • msf辅助模块详细操作
  • 【移动端网页特效】02-移动端轮播图(原生JS)
  • 神经网络(十二)卷积神经网络DLC
  • vue3.x 组件间传参
  • Tomcat域名访问文件出现访问不到的问题
  • BATJ 互联网公司面试必问知识点:Spring 全家桶全解,java 分布式框架技术方案
  • “大数据应用场景”之隔壁老王(连载四)
  • 4月23日世界读书日 网络营销论坛推荐《正在爆发的营销革命》
  • canvas实际项目操作,包含:线条,圆形,扇形,图片绘制,图片圆角遮罩,矩形,弧形文字...
  • HomeBrew常规使用教程
  • iOS动画编程-View动画[ 1 ] 基础View动画
  • Javascript Math对象和Date对象常用方法详解
  • JavaScript的使用你知道几种?(上)
  • spring-boot List转Page
  • Yeoman_Bower_Grunt
  • 阿里云Kubernetes容器服务上体验Knative
  • 阿里云应用高可用服务公测发布
  • 创建一个Struts2项目maven 方式
  • 开源地图数据可视化库——mapnik
  • 聊聊flink的TableFactory
  • 人脸识别最新开发经验demo
  • 如何借助 NoSQL 提高 JPA 应用性能
  • 使用iElevator.js模拟segmentfault的文章标题导航
  • puppet连载22:define用法
  • 整理一些计算机基础知识!
  • #Linux(帮助手册)
  • #经典论文 异质山坡的物理模型 2 有效导水率
  • #我与Java虚拟机的故事#连载16:打开Java世界大门的钥匙
  • $(selector).each()和$.each()的区别
  • (11)MSP430F5529 定时器B
  • (12)目标检测_SSD基于pytorch搭建代码
  • (32位汇编 五)mov/add/sub/and/or/xor/not
  • (Git) gitignore基础使用
  • (Matalb时序预测)WOA-BP鲸鱼算法优化BP神经网络的多维时序回归预测
  • (Redis使用系列) SpringBoot中Redis的RedisConfig 二
  • (翻译)Quartz官方教程——第一课:Quartz入门
  • (学习日记)2024.03.12:UCOSIII第十四节:时基列表
  • (一)Neo4j下载安装以及初次使用
  • (原)记一次CentOS7 磁盘空间大小异常的解决过程
  • (终章)[图像识别]13.OpenCV案例 自定义训练集分类器物体检测
  • (转)详解PHP处理密码的几种方式
  • (轉)JSON.stringify 语法实例讲解
  • ***汇编语言 实验16 编写包含多个功能子程序的中断例程
  • . Flume面试题
  • .NET 将多个程序集合并成单一程序集的 4+3 种方法
  • .net 使用$.ajax实现从前台调用后台方法(包含静态方法和非静态方法调用)