Skip to content

Commit 2df79a3

Browse files
alderneroaldernero
committed
added 2D kd-tree
1 parent 6b1a50d commit 2df79a3

File tree

5 files changed

+333
-0
lines changed

5 files changed

+333
-0
lines changed

‎geometry.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,14 @@ func (r *Rect) ContainsPoint(p Point) bool {
346346
return p.X >= r.X && p.X <= r.X+r.W && p.Y >= r.Y && p.Y <= r.Y+r.H
347347
}
348348

349+
func (r *Rect) Contains(rect Rect) bool {
350+
a := Point{X: r.X, Y: r.Y}
351+
b := Point{X: r.X + r.W, Y: r.Y + r.H}
352+
c := Point{X: rect.X, Y: rect.Y}
353+
d := Point{X: rect.X + rect.W, Y: rect.Y + rect.H}
354+
return a.X < c.X && a.Y < c.Y && b.X > d.X && b.Y > d.Y
355+
}
356+
349357
func (r *Rect) Intersects(rect Rect) bool {
350358
a := Point{X: r.X, Y: r.Y}
351359
b := Point{X: r.X + r.W, Y: r.Y + r.H}

‎kdtree.go

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
package sketchy
2+
3+
import "github.com/tdewolff/canvas"
4+
5+
type KDTree struct {
6+
point Point
7+
region Rect
8+
left *KDTree
9+
right *KDTree
10+
}
11+
12+
func NewKDTree(p Point, r Rect) *KDTree {
13+
return &KDTree{
14+
point: p,
15+
region: r,
16+
left: nil,
17+
right: nil,
18+
}
19+
}
20+
21+
func (k *KDTree) IsLeaf() bool {
22+
return k.left == nil && k.right == nil
23+
}
24+
25+
func (k *KDTree) Insert(p Point) {
26+
k.insert(p, 0)
27+
}
28+
29+
func (k *KDTree) Query(r Rect) []Point {
30+
var results []Point
31+
if k.IsLeaf() {
32+
if r.ContainsPoint(k.point) {
33+
results = append(results, k.point)
34+
}
35+
return results
36+
}
37+
if r.Contains(k.left.region) {
38+
results = append(results, k.left.reportSubtree()...)
39+
} else {
40+
if r.Intersects(k.left.region) {
41+
query := k.left.Query(r)
42+
if len(query) > 0 {
43+
results = append(results, query...)
44+
}
45+
}
46+
}
47+
if r.Contains(k.right.region) {
48+
results = append(results, k.right.reportSubtree()...)
49+
} else {
50+
if r.Intersects(k.right.region) {
51+
query := k.right.Query(r)
52+
if len(query) > 0 {
53+
results = append(results, query...)
54+
}
55+
}
56+
}
57+
return results
58+
}
59+
60+
func (k *KDTree) Size() int {
61+
count := 1
62+
if k.left != nil {
63+
count += k.left.Size()
64+
}
65+
if k.right != nil {
66+
count += k.right.Size()
67+
}
68+
return count
69+
}
70+
71+
func (k *KDTree) Draw(ctx *canvas.Context) {
72+
k.draw(ctx, 0, 0)
73+
}
74+
75+
func (k *KDTree) DrawWithPoints(s float64, ctx *canvas.Context) {
76+
k.draw(ctx, 0, s)
77+
}
78+
79+
func (k *KDTree) Clear() {
80+
k.left = nil
81+
k.right = nil
82+
}
83+
84+
func (k *KDTree) insert(p Point, d int) {
85+
if d%2 == 0 { // compare x value
86+
if p.X < k.point.X {
87+
if k.left == nil {
88+
rect := Rect{
89+
X: k.region.X,
90+
Y: k.region.Y,
91+
W: k.point.X - k.region.X,
92+
H: k.region.H,
93+
}
94+
k.left = NewKDTree(p, rect)
95+
} else {
96+
k.left.insert(p, d+1)
97+
}
98+
} else {
99+
if k.right == nil {
100+
rect := Rect{
101+
X: k.point.X,
102+
Y: k.region.Y,
103+
W: k.region.W - (k.point.X - k.region.X),
104+
H: k.region.H,
105+
}
106+
k.right = NewKDTree(p, rect)
107+
} else {
108+
k.right.insert(p, d+1)
109+
}
110+
}
111+
} else { // compare y value
112+
if p.Y < k.point.Y {
113+
if k.left == nil {
114+
rect := Rect{
115+
X: k.region.X,
116+
Y: k.region.Y,
117+
W: k.region.W,
118+
H: k.point.Y - k.region.Y,
119+
}
120+
k.left = NewKDTree(p, rect)
121+
} else {
122+
k.left.insert(p, d+1)
123+
}
124+
} else {
125+
if k.right == nil {
126+
rect := Rect{
127+
X: k.region.X,
128+
Y: k.point.Y,
129+
W: k.region.W,
130+
H: k.region.H - (k.point.Y - k.region.Y),
131+
}
132+
k.right = NewKDTree(p, rect)
133+
} else {
134+
k.right.insert(p, d+1)
135+
}
136+
}
137+
}
138+
}
139+
140+
func (k *KDTree) reportSubtree() []Point {
141+
var results []Point
142+
results = append(results, k.point)
143+
if k.left != nil {
144+
results = append(results, k.left.reportSubtree()...)
145+
}
146+
if k.right != nil {
147+
results = append(results, k.right.reportSubtree()...)
148+
}
149+
return results
150+
}
151+
152+
func (k *KDTree) draw(ctx *canvas.Context, depth int, pointSize float64) {
153+
if depth%2 == 0 {
154+
ctx.MoveTo(k.point.X, k.region.Y)
155+
ctx.LineTo(k.point.X, k.region.Y+k.region.H)
156+
ctx.Stroke()
157+
} else {
158+
ctx.MoveTo(k.region.X, k.point.Y)
159+
ctx.LineTo(k.region.X+k.region.W, k.point.Y)
160+
ctx.Stroke()
161+
}
162+
if pointSize > 0 {
163+
ctx.DrawPath(k.point.X, k.point.Y, canvas.Circle(pointSize))
164+
}
165+
166+
if k.left != nil {
167+
k.left.draw(ctx, depth+1, pointSize)
168+
}
169+
if k.right != nil {
170+
k.right.draw(ctx, depth+1, pointSize)
171+
}
172+
}

‎kdtree_test.go

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package sketchy
2+
3+
import (
4+
"github.com/stretchr/testify/assert"
5+
"math/rand"
6+
"testing"
7+
)
8+
9+
func TestKDTree_Insert(t *testing.T) {
10+
assert := assert.New(t)
11+
tree := NewKDTree(Point{X: 50, Y: 50}, Rect{X: 0, Y: 0, W: 100, H: 100})
12+
assert.Equal(1, tree.Size())
13+
}
14+
15+
func BenchmarkKDTree_Insert(b *testing.B) {
16+
var seed int64 = 42
17+
rand.Seed(seed)
18+
w := 210.0
19+
h := 297.0
20+
points := make([]Point, 1000)
21+
for i := 0; i < 1000; i++ {
22+
x := rand.Float64() * w
23+
y := rand.Float64() * h
24+
points[i] = Point{X: x, Y: y}
25+
}
26+
b.ResetTimer()
27+
for i := 0; i < b.N; i++ {
28+
tree := NewKDTree(points[0], Rect{X: 0, Y: 0, W: w, H: h})
29+
for j := 1; j < len(points); j++ {
30+
tree.Insert(points[j])
31+
}
32+
}
33+
}

‎visual_tests/kdtree_mouse/kdtree_mouse.go

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package main
2+
3+
import (
4+
"flag"
5+
"github.com/hajimehoshi/ebiten/v2/inpututil"
6+
"github.com/tdewolff/canvas"
7+
"image/color"
8+
"log"
9+
"math/rand"
10+
11+
"github.com/aldernero/sketchy"
12+
"github.com/hajimehoshi/ebiten/v2"
13+
)
14+
15+
var kdtree *sketchy.KDTree
16+
17+
func update(s *sketchy.Sketch) {
18+
// Update logic goes here
19+
if s.Toggle("Clear") {
20+
kdtree.Clear()
21+
}
22+
if inpututil.IsMouseButtonJustReleased(ebiten.MouseButtonLeft) {
23+
x, y := ebiten.CursorPosition()
24+
if s.PointInSketchArea(float64(x), float64(y)) {
25+
p := s.CanvasCoords(float64(x), float64(y))
26+
kdtree.Insert(p)
27+
}
28+
}
29+
}
30+
31+
func draw(s *sketchy.Sketch, c *canvas.Context) {
32+
// Drawing code goes here
33+
c.SetStrokeColor(color.White)
34+
c.SetFillColor(color.Transparent)
35+
c.SetStrokeCapper(canvas.ButtCap)
36+
c.SetStrokeWidth(s.Slider("Line Thickness"))
37+
if s.Toggle("Show Points") {
38+
kdtree.DrawWithPoints(s.Slider("Point Size"), c)
39+
} else {
40+
kdtree.Draw(c)
41+
}
42+
}
43+
44+
func main() {
45+
var configFile string
46+
var prefix string
47+
var randomSeed int64
48+
flag.StringVar(&configFile, "c", "sketch.json", "Sketch config file")
49+
flag.StringVar(&prefix, "p", "", "Output file prefix")
50+
flag.Int64Var(&randomSeed, "s", 0, "Random number generator seed")
51+
flag.Parse()
52+
s, err := sketchy.NewSketchFromFile(configFile)
53+
if err != nil {
54+
log.Fatal(err)
55+
}
56+
if prefix != "" {
57+
s.Prefix = prefix
58+
}
59+
s.RandomSeed = randomSeed
60+
s.Updater = update
61+
s.Drawer = draw
62+
s.Init()
63+
w := s.Width()
64+
h := s.Height()
65+
point := sketchy.Point{
66+
X: sketchy.Map(0, 1, 0.4*w, 0.6*w, rand.Float64()),
67+
Y: sketchy.Map(0, 1, 0.4*h, 0.6*h, rand.Float64()),
68+
}
69+
rect := sketchy.Rect{
70+
X: 0,
71+
Y: 0,
72+
W: w,
73+
H: h,
74+
}
75+
kdtree = sketchy.NewKDTree(point, rect)
76+
ebiten.SetWindowSize(int(s.ControlWidth+s.SketchWidth), int(s.SketchHeight))
77+
ebiten.SetWindowTitle("Sketchy - " + s.Title)
78+
ebiten.SetWindowResizable(false)
79+
ebiten.SetFPSMode(ebiten.FPSModeVsyncOffMaximum)
80+
ebiten.SetMaxTPS(ebiten.SyncWithFPS)
81+
if err := ebiten.RunGame(s); err != nil {
82+
log.Fatal(err)
83+
}
84+
}

‎visual_tests/kdtree_mouse/sketch.json

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
{
2+
"Title": "QuadTree Interaction Test",
3+
"SketchWidth": 800,
4+
"SketchHeight": 800,
5+
"ControlWidth": 240,
6+
"Sliders": [
7+
{
8+
"Name": "Line Thickness",
9+
"MinVal": 0.05,
10+
"MaxVal": 2,
11+
"Val": 0.3,
12+
"Incr": 0.05
13+
},
14+
{
15+
"Name": "Point Size",
16+
"MinVal": 0,
17+
"MaxVal": 5,
18+
"Val": 0.5,
19+
"Incr": 0.1
20+
}
21+
],
22+
"Toggles": [
23+
{
24+
"Name": "Show Points",
25+
"Checked": false
26+
},
27+
{
28+
"Name": "Clear",
29+
"IsButton": true
30+
}
31+
],
32+
"SketchBackgroundColor": "#1e1e1e",
33+
"SketchOutlineColor": "",
34+
"ControlBackgroundColor": "#1e1e1e",
35+
"ControlOutlineColor": "#ffdb00"
36+
}

0 commit comments

Comments
 (0)