Add RangeSearch, Add package kdrange

Author: Dennis Kuhnert

README.md

@@ -10,13 +10,11 @@
 A [k-d tree](https://en.wikipedia.org/wiki/K-d_tree) implementation in Go with:
 - n-dimensional points
 - k-nearest neighbor search
+- range search
 - remove without rebuilding the whole subtree
 - data attached to the points
 - using own structs by implementing a simple 2 function interface 
 
-Future features:
-- range search
-
 
 ## Usage
 
@@ -29,7 +27,6 @@ import "github.com/kyroy/kdtree"
 ````
 
 
-
 ### Implement the `kdtree.Point` interface
 
 ```go
@@ -64,12 +61,27 @@ func main() {
 	// KNN (k-nearest neighbor)
 	fmt.Println(tree.KNN(&points.Point{Coordinates: []float64{1, 1, 1}}, 2))
 	// [{3.00 1.00} {5.00 0.00}]
-
-	// other
-	fmt.Println(tree)
-	// [[[<nil> {3.00 1.00} {1.00 8.00}] {5.00 0.00} [<nil> {8.00 3.00} {7.00 5.00}]]]
+	
+	// RangeSearch
+	fmt.Println(tree.RangeSearch(kdrange.New(1, 8, 0, 2)))
+	// [{5.00 0.00} {3.00 1.00}]
+    
+	// Points
 	fmt.Println(tree.Points())
 	// [{3.00 1.00} {1.00 8.00} {5.00 0.00} {8.00 3.00} {7.00 5.00}]
+
+	// Remove
+	fmt.Println(tree.Remove(&points.Point2D{X: 5, Y: 0}))
+	// {5.00 0.00}
+
+	// String
+	fmt.Println(tree)
+	// [[{1.00 8.00} {3.00 1.00} [<nil> {8.00 3.00} {7.00 5.00}]]]
+
+	// Balance
+	tree.Balance()
+	fmt.Println(tree)
+	// [[[{3.00 1.00} {1.00 8.00} <nil>] {7.00 5.00} {8.00 3.00}]]
 }
 ```
 
@@ -81,9 +93,9 @@ type Data struct {
 
 func main() {
     tree := kdtree.New([]kdtree.Point{
-    points.NewPoint([]float64{7, 2, 3}, Data{value: "first"}),
-    points.NewPoint([]float64{3, 7, 10}, Data{value: "second"}),
-    points.NewPoint([]float64{4, 6, 1}, Data{value: "third"}),
+        points.NewPoint([]float64{7, 2, 3}, Data{value: "first"}),
+        points.NewPoint([]float64{3, 7, 10}, Data{value: "second"}),
+        points.NewPoint([]float64{4, 6, 1}, Data{value: "third"}),
     })
     
     // Insert
@@ -94,10 +106,25 @@ func main() {
     fmt.Println(tree.KNN(&points.Point{Coordinates: []float64{1, 1, 1}}, 2))
     // [{[4 6 1] {third}} {[7 2 3] {first}}]
     
-    // other
-    fmt.Println(tree)
-    // [[{[3 7 10] {second}} {[4 6 1] {third}} [{[8 1 0] {fifth}} {[7 2 3] {first}} {[12 4 6] {fourth}}]]]
+    // RangeSearch
+    fmt.Println(tree.RangeSearch(kdrange.New(1, 15, 1, 5, 0, 5)))
+    // [{[7 2 3] {first}} {[8 1 0] {fifth}}]
+    
+    // Points
     fmt.Println(tree.Points())
     // [{[3 7 10] {second}} {[4 6 1] {third}} {[8 1 0] {fifth}} {[7 2 3] {first}} {[12 4 6] {fourth}}]
+
+    // Remove
+    fmt.Println(tree.Remove(points.NewPoint([]float64{3, 7, 10}, nil)))
+    // {[3 7 10] {second}}
+
+    // String
+    fmt.Println(tree)
+    // [[<nil> {[4 6 1] {third}} [{[8 1 0] {fifth}} {[7 2 3] {first}} {[12 4 6] {fourth}}]]]
+
+    // Balance
+    tree.Balance()
+    fmt.Println(tree)
+    // [[[{[7 2 3] {first}} {[4 6 1] {third}} <nil>] {[8 1 0] {fifth}} {[12 4 6] {fourth}}]]
 }
 ```

kdrange/range.go

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2018 Dennis Kuhnert
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License");
+ *    you may not use this file except in compliance with the License.
+ *    You may obtain a copy of the License at
+ *
+ *        http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS,
+ *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *    See the License for the specific language governing permissions and
+ *    limitations under the License.
+ */
+
+// Package kdrange contains k-dimensional range struct and helpers.
+package kdrange
+
+// Range represents a range in k-dimensional space.
+type Range [][2]float64
+
+// New creates a new Range.
+//
+// It accepts a sequence of min/max pairs that define the Range.
+// For example a 2-dimensional rectangle with the with 2 and height 3, starting at (1,2):
+//
+//     r := NewRange(1, 3, 2, 5)
+//
+// I.e.:
+//     x (dim 0): 1 <= x <= 3
+//     y (dim 1): 2 <= y <= 5
+func New(limits ...float64) Range {
+	if limits == nil || len(limits)%2 != 0 {
+		return nil
+	}
+	r := make([][2]float64, len(limits)/2)
+	for i := range r {
+		r[i] = [2]float64{limits[2*i], limits[2*i+1]}
+	}
+	return r
+}

kdrange/range_test.go

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2018 Dennis Kuhnert
+ *
+ *    Licensed under the Apache License, Version 2.0 (the "License");
+ *    you may not use this file except in compliance with the License.
+ *    You may obtain a copy of the License at
+ *
+ *        http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *    Unless required by applicable law or agreed to in writing, software
+ *    distributed under the License is distributed on an "AS IS" BASIS,
+ *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *    See the License for the specific language governing permissions and
+ *    limitations under the License.
+ */
+
+package kdrange_test
+
+import (
+	"testing"
+
+	"github.com/kyroy/kdtree/kdrange"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewRange(t *testing.T) {
+	tests := []struct {
+		name   string
+		input  []float64
+		output kdrange.Range
+	}{
+		{
+			name:   "nil",
+			input:  nil,
+			output: nil,
+		},
+		{
+			name:   "uneven",
+			input:  []float64{0},
+			output: nil,
+		},
+		{
+			name:   "empty",
+			input:  []float64{},
+			output: [][2]float64{},
+		},
+		{
+			name:   "2d",
+			input:  []float64{1, 2, 3, 4},
+			output: [][2]float64{{1, 2}, {3, 4}},
+		},
+		{
+			name:   "3d",
+			input:  []float64{4, 1, 6, 19, 3, 1},
+			output: [][2]float64{{4, 1}, {6, 19}, {3, 1}},
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			assert.Equal(t, test.output, kdrange.New(test.input...))
+		})
+	}
+}

kdtree.go

@@ -22,6 +22,7 @@ import (
 	"math"
 	"sort"
 
+	"github.com/kyroy/kdtree/kdrange"
 	"github.com/kyroy/priority-queue"
 )
 
@@ -116,6 +117,7 @@ func (t *KDTree) Points() []Point {
 }
 
 // KNN returns the k-nearest neighbours of the given point.
+// The points are sorted by the distance to the given points. Starting with the nearest.
 func (t *KDTree) KNN(p Point, k int) []Point {
 	if t.root == nil || p == nil || k == 0 {
 		return []Point{}
@@ -133,10 +135,22 @@ func (t *KDTree) KNN(p Point, k int) []Point {
 	return points
 }
 
+// RangeSearch returns all points in the given range r.
+//
+// Returns an empty slice when input is nil or len(r) does not equal Point.Dimensions().
+func (t *KDTree) RangeSearch(r kdrange.Range) []Point {
+	if t.root == nil || r == nil || len(r) != t.root.Dimensions() {
+		return []Point{}
+	}
+
+	return t.root.RangeSearch(r, 0)
+}
+
 func knn(p Point, k int, start *node, currentAxis int, nearestPQ *pq.PriorityQueue) {
 	if p == nil || k == 0 || start == nil {
 		return
 	}
+
 	var path []*node
 	currentNode := start
 
@@ -353,3 +367,24 @@ func (n *node) FindLargest(axis int, largest *node) *node {
 	}
 	return largest
 }
+
+func (n *node) RangeSearch(r kdrange.Range, axis int) []Point {
+	points := []Point{}
+
+	for dim, limit := range r {
+		if limit[0] > n.Dimension(dim) || limit[1] < n.Dimension(dim) {
+			goto checkChildren
+		}
+	}
+	points = append(points, n.Point)
+
+checkChildren:
+	if n.Left != nil && n.Dimension(axis) >= r[axis][0] {
+		points = append(points, n.Left.RangeSearch(r, (axis+1)%n.Dimensions())...)
+	}
+	if n.Right != nil && n.Dimension(axis) <= r[axis][1] {
+		points = append(points, n.Right.RangeSearch(r, (axis+1)%n.Dimensions())...)
+	}
+
+	return points
+}