1 | ;; |
---|
2 | ;; An implementation of the K-D tree spatial indexing data structures. |
---|
3 | ;; http://en.wikipedia.org/wiki/K-d_tree |
---|
4 | ;; |
---|
5 | ;; This code is based on the Haskell kd-tree library implementation of |
---|
6 | ;; K-D trees. |
---|
7 | ;; |
---|
8 | ;; Copyright 2012 Ivan Raikov and the Okinawa Institute of |
---|
9 | ;; Science and Technology. |
---|
10 | ;; |
---|
11 | ;; This program is free software: you can redistribute it and/or |
---|
12 | ;; modify it under the terms of the GNU General Public License as |
---|
13 | ;; published by the Free Software Foundation, either version 3 of the |
---|
14 | ;; License, or (at your option) any later version. |
---|
15 | ;; |
---|
16 | ;; This program is distributed in the hope that it will be useful, but |
---|
17 | ;; WITHOUT ANY WARRANTY; without even the implied warranty of |
---|
18 | ;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
---|
19 | ;; General Public License for more details. |
---|
20 | ;; |
---|
21 | ;; A full copy of the GPL license can be found at |
---|
22 | ;; <http://www.gnu.org/licenses/>. |
---|
23 | ;; |
---|
24 | |
---|
25 | (module kd-tree |
---|
26 | |
---|
27 | ( |
---|
28 | <Point> default-<Point> Point3d Point2d |
---|
29 | point? make-point |
---|
30 | |
---|
31 | <KdTree> default-<KdTree> KdTree3d KdTree2d |
---|
32 | |
---|
33 | kd-tree? |
---|
34 | kd-tree-empty? |
---|
35 | kd-tree->list |
---|
36 | kd-tree->list* |
---|
37 | kd-tree-map |
---|
38 | kd-tree-for-each |
---|
39 | kd-tree-for-each* |
---|
40 | kd-tree-fold-right |
---|
41 | kd-tree-fold-right* |
---|
42 | kd-tree-subtrees |
---|
43 | kd-tree-node-points |
---|
44 | kd-tree-node-indices |
---|
45 | kd-tree-min-index |
---|
46 | kd-tree-max-index |
---|
47 | kd-tree-size |
---|
48 | ) |
---|
49 | |
---|
50 | (import scheme chicken data-structures foreign) |
---|
51 | |
---|
52 | (require-library srfi-1 srfi-4 extras cis) |
---|
53 | (require-extension typeclass datatype) |
---|
54 | |
---|
55 | (import (only srfi-1 xcons fold list-tabulate split-at every fold-right take filter filter-map remove zip) |
---|
56 | (only srfi-4 f64vector-ref f64vector-set! f64vector-length make-f64vector f64vector->list |
---|
57 | u32vector-ref u32vector-set!) |
---|
58 | (only extras fprintf pp) |
---|
59 | (only foreign foreign-lambda) |
---|
60 | (prefix cis cis:)) |
---|
61 | |
---|
62 | |
---|
63 | (define-class <Point> |
---|
64 | ;; dimension returns the number of coordinates of a point. |
---|
65 | dimension ;; Point -> Int |
---|
66 | |
---|
67 | ;; gets the k'th coordinate, starting from 0. |
---|
68 | coord ;; Int * Point -> Double |
---|
69 | |
---|
70 | ;; compares the given coordinates |
---|
71 | compare-coord ;; Int * Point * Point -> Bool |
---|
72 | |
---|
73 | ;; returns the squared distance between two points. |
---|
74 | dist2 ;; Point * Point -> Double |
---|
75 | |
---|
76 | ;; returns the scaled squared distance between two points. |
---|
77 | sdist2 ;; Point * Point * [Int] -> Double |
---|
78 | |
---|
79 | ;; returns 0, negative or positive number depending on the |
---|
80 | ;; distance between two points |
---|
81 | compare-distance |
---|
82 | |
---|
83 | ) |
---|
84 | |
---|
85 | (define (minimum-by lst less? . rest) |
---|
86 | (if (null? lst) #f |
---|
87 | (if (null? rest) |
---|
88 | (let recur ((lst (cdr lst)) (m (car lst))) |
---|
89 | (if (null? lst) m |
---|
90 | (if (less? (car lst) m) |
---|
91 | (recur (cdr lst) (car lst)) |
---|
92 | (recur (cdr lst) m) |
---|
93 | )) |
---|
94 | ) |
---|
95 | (let recur ((lst (cdr lst)) |
---|
96 | (rest (map cdr rest)) |
---|
97 | (m (map car (cons lst rest)))) |
---|
98 | (if (null? lst) m |
---|
99 | (if (less? (car lst) (car m)) |
---|
100 | (recur (cdr lst) (map cdr rest) (map car (cons lst rest))) |
---|
101 | (recur (cdr lst) (map cdr rest) m) |
---|
102 | )) |
---|
103 | ) |
---|
104 | ))) |
---|
105 | |
---|
106 | (define (sum lst) (fold + 0. lst)) |
---|
107 | |
---|
108 | (define (default-<Point> dimension coord) |
---|
109 | |
---|
110 | (let* ((dist2 |
---|
111 | (lambda (a b) |
---|
112 | (let ((diff2 (lambda (i) (let ((v (- (coord i a) (coord i b)))) (* v v))))) |
---|
113 | (sum (list-tabulate (dimension a) diff2))))) |
---|
114 | |
---|
115 | (sdist2 |
---|
116 | (lambda (factors) |
---|
117 | (let ((factors2 (map (lambda (n) (* n n)) factors))) |
---|
118 | (lambda (a b) |
---|
119 | (let ((diff2 (lambda (i) (let ((v (- (coord i a) (coord i b)))) (* v v))))) |
---|
120 | (let ((v (sum (map * (list-tabulate (dimension a) diff2) factors2)))) |
---|
121 | v)))))) |
---|
122 | |
---|
123 | (compare-distance |
---|
124 | (lambda (p a b . reltol) |
---|
125 | (let ((delta (- (dist2 p a) (dist2 p b)))) |
---|
126 | (if (null? reltol) |
---|
127 | delta |
---|
128 | (if (<= delta (car reltol)) 0 delta))))) |
---|
129 | |
---|
130 | (compare-coord |
---|
131 | (lambda (c a b) |
---|
132 | (< (coord c a) (coord c b)))) |
---|
133 | |
---|
134 | ) |
---|
135 | |
---|
136 | (make-<Point> |
---|
137 | dimension coord compare-coord dist2 sdist2 compare-distance) |
---|
138 | )) |
---|
139 | |
---|
140 | (define point? vector?) |
---|
141 | (define make-point vector) |
---|
142 | |
---|
143 | |
---|
144 | (define Point3d |
---|
145 | (default-<Point> |
---|
146 | (lambda (p) (and (point? p) 3)) |
---|
147 | (lambda (i p) (vector-ref p i)) |
---|
148 | )) |
---|
149 | |
---|
150 | (define Point2d |
---|
151 | (default-<Point> |
---|
152 | (lambda (p) (and (point? p) 2)) |
---|
153 | (lambda (i p) (vector-ref p i)) |
---|
154 | )) |
---|
155 | |
---|
156 | |
---|
157 | (define-class <KdTree> |
---|
158 | |
---|
159 | ;; constructs a kd-tree from a list of points |
---|
160 | list->kd-tree |
---|
161 | ;; nearest neighbor of a point |
---|
162 | kd-tree-nearest-neighbor |
---|
163 | ;; the index of the nearest neighbor of a point |
---|
164 | kd-tree-nearest-neighbor* |
---|
165 | ;; neighbors of a point within radius r |
---|
166 | kd-tree-near-neighbors |
---|
167 | ;; neighbors of a point within radius r (using point indices) |
---|
168 | kd-tree-near-neighbors* |
---|
169 | ;; k nearest neighbors of a point |
---|
170 | kd-tree-k-nearest-neighbors |
---|
171 | ;; removes a point from the tree |
---|
172 | kd-tree-remove |
---|
173 | ;; retrieves all points between two planes |
---|
174 | kd-tree-slice |
---|
175 | ;; retrieves all points between two planes (using point indices) |
---|
176 | kd-tree-slice* |
---|
177 | ;; checks that the kd-tree properties are preserved |
---|
178 | kd-tree-is-valid? |
---|
179 | kd-tree-all-subtrees-are-valid? |
---|
180 | |
---|
181 | ) |
---|
182 | |
---|
183 | |
---|
184 | (define (positive-or-zero-integer? x) |
---|
185 | (and (integer? x) (or (zero? x) (positive? x)))) |
---|
186 | |
---|
187 | |
---|
188 | (define (positive-integer? x) |
---|
189 | (and (integer? x) (positive? x))) |
---|
190 | |
---|
191 | |
---|
192 | (define-datatype kd-tree kd-tree? |
---|
193 | (KdNode (left kd-tree?) |
---|
194 | (p point?) |
---|
195 | (i positive-or-zero-integer?) |
---|
196 | (v (lambda (v) (or (not v) v))) |
---|
197 | (right kd-tree?) |
---|
198 | (axis positive-or-zero-integer?) |
---|
199 | (ci cis:cis?) |
---|
200 | ) |
---|
201 | (KdLeaf (ii cis:cis?) |
---|
202 | (pp (lambda (lst) (every point? lst))) |
---|
203 | (vv (lambda (v) (or (list? v) (not v)))) |
---|
204 | (axis positive-or-zero-integer?) ) |
---|
205 | ) |
---|
206 | |
---|
207 | |
---|
208 | (define (kd-tree-empty? t) |
---|
209 | (cases kd-tree t |
---|
210 | (KdLeaf (ii pp vv axis) (cis:empty? ii)) |
---|
211 | (else #f))) |
---|
212 | |
---|
213 | |
---|
214 | (define (kd-tree-map f t) |
---|
215 | (cases kd-tree t |
---|
216 | (KdLeaf (ii pp vv axis) |
---|
217 | (KdLeaf ii (map f pp) vv axis)) |
---|
218 | (KdNode (l x i v r axis ci) |
---|
219 | (KdNode (kd-tree-map f l) |
---|
220 | (f x) |
---|
221 | i v |
---|
222 | (kd-tree-map f r) |
---|
223 | axis ci)) |
---|
224 | )) |
---|
225 | |
---|
226 | (define (kd-tree-for-each f t) |
---|
227 | (cases kd-tree t |
---|
228 | (KdLeaf (ii pp vv axis) (for-each f pp)) |
---|
229 | (KdNode (l x i v r axis ci) |
---|
230 | (begin |
---|
231 | (kd-tree-for-each f l) |
---|
232 | (f x) |
---|
233 | (kd-tree-for-each f r) |
---|
234 | )) |
---|
235 | )) |
---|
236 | |
---|
237 | |
---|
238 | (define (kd-tree-for-each* f t) |
---|
239 | (cases kd-tree t |
---|
240 | (KdLeaf (ii pp vv axis) |
---|
241 | (if vv (for-each f (zip (reverse (cis:elements ii)) vv) pp) |
---|
242 | (for-each f (reverse (cis:elements ii)) pp))) |
---|
243 | (KdNode (l x i v r axis ci) |
---|
244 | (begin |
---|
245 | (kd-tree-for-each* f l) |
---|
246 | (if v (f (list i v) x) (f i x)) |
---|
247 | (kd-tree-for-each* f r) |
---|
248 | )) |
---|
249 | )) |
---|
250 | |
---|
251 | |
---|
252 | (define (kd-tree-fold-right f init t) |
---|
253 | (cases kd-tree t |
---|
254 | (KdLeaf (ii pp vv axis) |
---|
255 | (fold-right f init pp)) |
---|
256 | (KdNode (l p i v r axis ci) |
---|
257 | (let* ((init2 (kd-tree-fold-right f init r)) |
---|
258 | (init3 (f p init2))) |
---|
259 | (kd-tree-fold-right f init3 l))) |
---|
260 | )) |
---|
261 | |
---|
262 | |
---|
263 | (define (kd-tree-fold-right* f init t) |
---|
264 | (cases kd-tree t |
---|
265 | (KdLeaf (ii pp vv axis) |
---|
266 | (if vv |
---|
267 | (fold-right f init (zip (reverse (cis:elements ii)) vv) pp) |
---|
268 | (fold-right f init (reverse (cis:elements ii)) pp))) |
---|
269 | (KdNode (l x i v r axis ci) |
---|
270 | (let* ((init2 (kd-tree-fold-right* f init r)) |
---|
271 | (init3 (if v (f (list i v) x init2) (f i x init2)))) |
---|
272 | (kd-tree-fold-right* f init3 l))) |
---|
273 | )) |
---|
274 | |
---|
275 | |
---|
276 | |
---|
277 | |
---|
278 | (define (kd-tree->list t) |
---|
279 | (kd-tree-fold-right cons '() t)) |
---|
280 | |
---|
281 | |
---|
282 | (define (kd-tree->list* t) |
---|
283 | (kd-tree-fold-right* |
---|
284 | (lambda (iv x ax) (cons (list iv x) ax)) |
---|
285 | '() t)) |
---|
286 | |
---|
287 | |
---|
288 | ;; Returns a list containing t and all its subtrees, including the |
---|
289 | ;; leaf nodes. |
---|
290 | |
---|
291 | (define (kd-tree-subtrees t) |
---|
292 | (cases kd-tree t |
---|
293 | (KdLeaf (ii pp vv axis) (list t)) |
---|
294 | (KdNode (l x i v r axis ci) |
---|
295 | (append (kd-tree-subtrees l) |
---|
296 | (list t) |
---|
297 | (kd-tree-subtrees r))) |
---|
298 | )) |
---|
299 | |
---|
300 | |
---|
301 | (define (kd-tree-node-points t) |
---|
302 | (cases kd-tree t |
---|
303 | (KdLeaf (ii pp vv axis) pp) |
---|
304 | (KdNode (l x i v r axis ci) (list x)) |
---|
305 | )) |
---|
306 | |
---|
307 | (define (kd-tree-node-indices t) |
---|
308 | (cases kd-tree t |
---|
309 | (KdLeaf (ii pp vv axis) (cis:elements ii)) |
---|
310 | (KdNode (l x i v r axis ci) (list i)) |
---|
311 | )) |
---|
312 | |
---|
313 | (define (kd-tree-node-values t) |
---|
314 | (cases kd-tree t |
---|
315 | (KdLeaf (ii pp vv axis) vv) |
---|
316 | (KdNode (l x i v r axis ci) (list v)) |
---|
317 | )) |
---|
318 | |
---|
319 | (define (kd-tree-size t) |
---|
320 | (cases kd-tree t |
---|
321 | (KdLeaf (ii pp vv axis) (cis:cardinal ii)) |
---|
322 | (KdNode (l x i v r axis ci) (cis:cardinal ci)))) |
---|
323 | |
---|
324 | (define (kd-tree-min-index t) |
---|
325 | (cases kd-tree t |
---|
326 | (KdLeaf (ii pp vv axis) (cis:get-min ii)) |
---|
327 | (KdNode (l x i v r axis ci) (cis:get-min ci)))) |
---|
328 | |
---|
329 | (define (kd-tree-max-index t) |
---|
330 | (cases kd-tree t |
---|
331 | (KdLeaf (ii pp vv axis) (cis:get-max ii)) |
---|
332 | (KdNode (l x i v r axis ci) (cis:get-max ci)))) |
---|
333 | |
---|
334 | |
---|
335 | |
---|
336 | |
---|
337 | ;; construct a kd-tree from a list of points |
---|
338 | (define=> (make-list->kd-tree/depth <Point>) |
---|
339 | |
---|
340 | (lambda (make-point make-value) |
---|
341 | |
---|
342 | (letrec ( |
---|
343 | (split |
---|
344 | (lambda (m n points depth) |
---|
345 | |
---|
346 | (let* ((axis (modulo depth (dimension (make-point (car points))))) |
---|
347 | (cmpfn (lambda (p0 p1) |
---|
348 | (compare-coord axis (make-point p0) (make-point p1)))) |
---|
349 | (sorted (sort points cmpfn)) |
---|
350 | (median-index (quotient (- n m) 2))) |
---|
351 | |
---|
352 | (let-values (((lt gte) (split-at sorted median-index))) |
---|
353 | (values (car gte) median-index lt (cdr gte)))) |
---|
354 | )) |
---|
355 | |
---|
356 | (list->kd-tree/depth |
---|
357 | (lambda (m n points depth #!key (leaf-factor 10) (offset 0)) |
---|
358 | |
---|
359 | (cond |
---|
360 | ((null? points) (KdLeaf cis:empty '() '() depth)) |
---|
361 | |
---|
362 | ((<= (- n m) leaf-factor) |
---|
363 | (let ((k (- n m))) |
---|
364 | |
---|
365 | (let* ((es (take points k)) |
---|
366 | (ps (map make-point es)) |
---|
367 | (ii (cis:shift offset (cis:interval m (- n 1)))) |
---|
368 | (vs (and make-value |
---|
369 | (map (lambda (i e) (make-value i e)) |
---|
370 | (reverse (cis:elements ii)) |
---|
371 | es))) |
---|
372 | ) |
---|
373 | |
---|
374 | (KdLeaf ii ps vs (modulo depth (dimension (car ps)))) |
---|
375 | ))) |
---|
376 | |
---|
377 | ((null? (cdr points)) |
---|
378 | (let* ((e (car points)) |
---|
379 | (ps (list (make-point e)) ) |
---|
380 | (vs (and make-value (list (make-value m e))))) |
---|
381 | |
---|
382 | (KdLeaf (cis:shift offset (cis:singleton m) ) |
---|
383 | ps vs |
---|
384 | (modulo depth (dimension (car ps)))) |
---|
385 | )) |
---|
386 | |
---|
387 | (else |
---|
388 | (let-values (((median median-index lt gte) |
---|
389 | (split m n points depth))) |
---|
390 | |
---|
391 | |
---|
392 | (let* ((depth1 (+ 1 depth)) |
---|
393 | (i (+ m median-index offset)) |
---|
394 | (p (make-point median)) |
---|
395 | (v (and make-value (make-value i median))) |
---|
396 | |
---|
397 | (axis (modulo depth (dimension p)))) |
---|
398 | |
---|
399 | (KdNode (list->kd-tree/depth m (+ m median-index) lt depth1 |
---|
400 | leaf-factor: leaf-factor) |
---|
401 | p i v |
---|
402 | (list->kd-tree/depth (+ m median-index 1) n gte depth1 |
---|
403 | leaf-factor: leaf-factor) |
---|
404 | axis |
---|
405 | (cis:shift offset (cis:interval m (- n 1)))))) |
---|
406 | ))) |
---|
407 | )) |
---|
408 | list->kd-tree/depth |
---|
409 | ))) |
---|
410 | |
---|
411 | ;; Returns the nearest neighbor of p in tree t. |
---|
412 | |
---|
413 | (define=> (make-kd-tree-nearest-neighbor <Point>) |
---|
414 | (define (tree-empty? t) (cases kd-tree t (KdLeaf (ii pp vv axis) (cis:empty? ii)) (else #f))) |
---|
415 | (letrec ((find-nearest |
---|
416 | (lambda (t1 t2 p probe xp x-probe) |
---|
417 | |
---|
418 | (let* ((candidates1 |
---|
419 | (let ((best1 (nearest-neighbor t1 probe))) |
---|
420 | (or (and best1 (list best1 p)) (list p)))) |
---|
421 | |
---|
422 | (sphere-intersects-plane? |
---|
423 | (let ((v (- x-probe xp))) |
---|
424 | (< (* v v) (dist2 probe (car candidates1))))) |
---|
425 | |
---|
426 | (candidates2 |
---|
427 | (if sphere-intersects-plane? |
---|
428 | (let ((nn (nearest-neighbor t2 probe))) |
---|
429 | (if nn (append candidates1 (list nn)) candidates1)) |
---|
430 | candidates1))) |
---|
431 | |
---|
432 | (minimum-by candidates2 (lambda (a b) (negative? (compare-distance probe a b)))) |
---|
433 | ))) |
---|
434 | |
---|
435 | |
---|
436 | (nearest-neighbor |
---|
437 | (lambda (t probe) |
---|
438 | (cases kd-tree t |
---|
439 | (KdLeaf (ii pp vv axis) |
---|
440 | (minimum-by pp (lambda (a b) (negative? (compare-distance probe a b))))) |
---|
441 | |
---|
442 | (KdNode (l p i v r axis ci) |
---|
443 | (if (and (tree-empty? l) |
---|
444 | (tree-empty? r)) p |
---|
445 | (let ((x-probe (coord axis probe)) |
---|
446 | (xp (coord axis p))) |
---|
447 | |
---|
448 | (if (<= x-probe xp) |
---|
449 | (find-nearest l r p probe xp x-probe) |
---|
450 | (find-nearest r l p probe xp x-probe)) |
---|
451 | )) |
---|
452 | )) |
---|
453 | ))) |
---|
454 | |
---|
455 | nearest-neighbor |
---|
456 | )) |
---|
457 | |
---|
458 | ;; Returns the index of the nearest neighbor of p in tree t. |
---|
459 | |
---|
460 | (define=> (make-kd-tree-nearest-neighbor* <Point>) |
---|
461 | |
---|
462 | (define (tree-empty? t) (cases kd-tree t (KdLeaf (ii pp vv axis) (cis:empty? ii)) (else #f))) |
---|
463 | |
---|
464 | (letrec ((find-nearest |
---|
465 | (lambda (t1 t2 i v p probe xp x-probe) |
---|
466 | |
---|
467 | (let* ((candidates1 |
---|
468 | (let ((best1 (nearest-neighbor t1 probe))) |
---|
469 | (or (and best1 (list best1 (if v (list (list i v) p) (list i p)))) |
---|
470 | (list (if v (list (list i v) p) (list i p)))) |
---|
471 | )) |
---|
472 | |
---|
473 | (sphere-intersects-plane? |
---|
474 | (let ((v (- x-probe xp))) |
---|
475 | (< (* v v) (dist2 probe (cadar candidates1))))) |
---|
476 | |
---|
477 | (candidates2 |
---|
478 | (if sphere-intersects-plane? |
---|
479 | (let ((nn (nearest-neighbor t2 probe))) |
---|
480 | (if nn (append candidates1 (list nn)) candidates1)) |
---|
481 | candidates1))) |
---|
482 | |
---|
483 | (let ((v (minimum-by candidates2 (lambda (a b) (negative? (compare-distance probe (cadr a) (cadr b))))))) |
---|
484 | v) |
---|
485 | ))) |
---|
486 | |
---|
487 | |
---|
488 | (nearest-neighbor |
---|
489 | (lambda (t probe) |
---|
490 | (cases kd-tree t |
---|
491 | (KdLeaf (ii pp vv axis) |
---|
492 | (let ((res |
---|
493 | (if vv |
---|
494 | (minimum-by pp (lambda (a b) (negative? (compare-distance probe a b))) |
---|
495 | (zip (reverse (cis:elements ii)) vv )) |
---|
496 | (minimum-by pp (lambda (a b) (negative? (compare-distance probe a b))) |
---|
497 | (reverse (cis:elements ii)))))) |
---|
498 | (and res (reverse res)))) |
---|
499 | |
---|
500 | (KdNode (l p i v r axis ci) |
---|
501 | |
---|
502 | (if (and (tree-empty? l) |
---|
503 | (tree-empty? r)) |
---|
504 | |
---|
505 | (if v (list (list i v) p) (list i p)) |
---|
506 | |
---|
507 | (let ((x-probe (coord axis probe)) |
---|
508 | (xp (coord axis p)) |
---|
509 | (xi i)) |
---|
510 | (if (<= x-probe xp) |
---|
511 | (find-nearest l r i v p probe xp x-probe) |
---|
512 | (find-nearest r l i v p probe xp x-probe)) |
---|
513 | )) |
---|
514 | )) |
---|
515 | ))) |
---|
516 | |
---|
517 | nearest-neighbor |
---|
518 | )) |
---|
519 | |
---|
520 | |
---|
521 | ;; near-neighbors t r p returns all neighbors within distance r from p in tree t. |
---|
522 | |
---|
523 | (define=> (make-kd-tree-near-neighbors <Point>) |
---|
524 | (define (tree-empty? t) (cases kd-tree t (KdLeaf (ii pp vv axis) (cis:empty? ii)) (else #f))) |
---|
525 | (letrec ((near-neighbors |
---|
526 | (lambda (t radius probe fdist filter-fn) |
---|
527 | (cases kd-tree t |
---|
528 | (KdLeaf (ii pp vv axis) |
---|
529 | (let ((r2 (* radius radius))) |
---|
530 | (filter-fn probe pp r2))) |
---|
531 | |
---|
532 | (KdNode (l p i v r axis ci) |
---|
533 | (let ((maybe-pivot (filter-fn probe (list p) (* radius radius)))) |
---|
534 | |
---|
535 | (if (and (tree-empty? l) |
---|
536 | (tree-empty? r)) |
---|
537 | |
---|
538 | maybe-pivot |
---|
539 | |
---|
540 | (let ((x-probe (coord axis probe)) |
---|
541 | (xp (coord axis p))) |
---|
542 | |
---|
543 | (if (<= x-probe xp) |
---|
544 | |
---|
545 | (let ((nearest (append maybe-pivot (near-neighbors l radius probe fdist filter-fn)))) |
---|
546 | (if (> (+ x-probe (abs radius)) xp) |
---|
547 | (append (near-neighbors r radius probe fdist filter-fn) nearest) |
---|
548 | nearest)) |
---|
549 | |
---|
550 | (let ((nearest (append maybe-pivot (near-neighbors r radius probe fdist filter-fn)))) |
---|
551 | (if (< (- x-probe (abs radius)) xp) |
---|
552 | (append (near-neighbors l radius probe fdist filter-fn) nearest) |
---|
553 | nearest))) |
---|
554 | )))) |
---|
555 | )) |
---|
556 | )) |
---|
557 | (lambda (t radius probe #!key (factors #f) (with-distance? #f)) |
---|
558 | (let* ((dist-fn (if factors (sdist2 factors) dist2)) |
---|
559 | (filter-fn (if with-distance? |
---|
560 | (lambda (probe pp d2) (filter-map (lambda (p) (let ((pd (dist-fn probe p))) |
---|
561 | (and (<= pd d2) (list p (sqrt pd) )))) pp)) |
---|
562 | (lambda (probe pp d2) (filter (lambda (p) (<= (dist-fn probe p) d2)) pp)) |
---|
563 | )) |
---|
564 | ) |
---|
565 | |
---|
566 | (near-neighbors t radius probe dist-fn filter-fn) |
---|
567 | )) |
---|
568 | )) |
---|
569 | |
---|
570 | |
---|
571 | (define=> (make-kd-tree-near-neighbors* <Point>) |
---|
572 | (define (tree-empty? t) (cases kd-tree t (KdLeaf (ii pp vv axis) (cis:empty? ii)) (else #f))) |
---|
573 | (letrec ((near-neighbors |
---|
574 | (lambda (t radius probe fdist filter-fn) |
---|
575 | (cases kd-tree t |
---|
576 | |
---|
577 | (KdLeaf (ii pp vv axis) |
---|
578 | (let ((r2 (* radius radius))) |
---|
579 | (filter-fn probe pp ii vv r2))) |
---|
580 | |
---|
581 | (KdNode (l p i v r axis ci) |
---|
582 | (let ((maybe-pivot (filter-fn probe (list p) (list i) (list v) (* radius radius)))) |
---|
583 | |
---|
584 | (if (and (tree-empty? l) |
---|
585 | (tree-empty? r)) |
---|
586 | |
---|
587 | maybe-pivot |
---|
588 | |
---|
589 | (let ((x-probe (coord axis probe)) |
---|
590 | (xp (coord axis p))) |
---|
591 | |
---|
592 | (if (<= x-probe xp) |
---|
593 | |
---|
594 | (let ((nearest (append maybe-pivot (near-neighbors l radius probe fdist filter-fn)))) |
---|
595 | (if (> (+ x-probe (abs radius)) xp) |
---|
596 | (append (near-neighbors r radius probe fdist filter-fn) nearest) |
---|
597 | nearest)) |
---|
598 | |
---|
599 | (let ((nearest (append maybe-pivot (near-neighbors r radius probe fdist filter-fn)))) |
---|
600 | (if (< (- x-probe (abs radius)) xp) |
---|
601 | (append (near-neighbors l radius probe fdist filter-fn) nearest) |
---|
602 | nearest))) |
---|
603 | )) |
---|
604 | )) |
---|
605 | )) |
---|
606 | )) |
---|
607 | (lambda (t radius probe #!key (factors #f) (with-distance? #f)) |
---|
608 | (let* ((dist-fn (if factors (sdist2 factors) dist2)) |
---|
609 | (filter-fn (if with-distance? |
---|
610 | (lambda (probe pp ii vv r2) |
---|
611 | (if vv |
---|
612 | (filter-map (lambda (i v p) (let ((pd (dist-fn probe p))) (and (<= pd r2) (list (list i v) p (sqrt pd))))) |
---|
613 | (reverse (cis:elements ii)) vv pp) |
---|
614 | (filter-map (lambda (i p) (let ((pd (dist-fn probe p))) (and (<= pd r2) (list i p (sqrt pd))))) |
---|
615 | (reverse (cis:elements ii)) pp) |
---|
616 | )) |
---|
617 | (lambda (probe pp ii vv r2) |
---|
618 | (if vv |
---|
619 | (filter-map (lambda (i v p) (and (<= (dist-fn probe p) r2) (list (list i v) p))) |
---|
620 | (reverse (cis:elements ii)) vv pp) |
---|
621 | (filter-map (lambda (i p) (and (<= (dist-fn probe p) r2) (list i p))) |
---|
622 | (reverse (cis:elements ii)) pp) |
---|
623 | )) |
---|
624 | )) |
---|
625 | ) |
---|
626 | |
---|
627 | (near-neighbors t radius probe dist-fn filter-fn) |
---|
628 | |
---|
629 | )) |
---|
630 | )) |
---|
631 | |
---|
632 | |
---|
633 | |
---|
634 | ;; Returns the k nearest points to p within tree. |
---|
635 | (define=> (make-kd-tree-k-nearest-neighbors <Point>) |
---|
636 | (lambda (kd-tree-remove kd-tree-nearest-neighbor) |
---|
637 | (letrec ((k-nearest-neighbors |
---|
638 | (lambda (t k probe) |
---|
639 | (cases kd-tree t |
---|
640 | |
---|
641 | (KdLeaf (ii pp vv axis) |
---|
642 | (let recur ((res '()) (pp pp) (k k)) |
---|
643 | (if (or (<= k 0) (null? pp)) |
---|
644 | res |
---|
645 | (let ((nearest (minimum-by pp (lambda (a b) (negative? (compare-distance probe a b)))))) |
---|
646 | (recur (cons nearest res) |
---|
647 | (remove (lambda (p) (equal? p nearest)) pp) |
---|
648 | (- k 1)) |
---|
649 | )) |
---|
650 | )) |
---|
651 | |
---|
652 | (else |
---|
653 | (if (<= k 0) '() |
---|
654 | (let* ((nearest (kd-tree-nearest-neighbor t probe)) |
---|
655 | (tree1 (kd-tree-remove t nearest))) |
---|
656 | (cons nearest (k-nearest-neighbors tree1 (- k 1) probe))) |
---|
657 | )) |
---|
658 | )) |
---|
659 | )) |
---|
660 | k-nearest-neighbors))) |
---|
661 | |
---|
662 | |
---|
663 | ;; removes the point p from t. |
---|
664 | (define=> (make-kd-tree-remove <Point>) |
---|
665 | (lambda (list->kd-tree list->kd-tree*) |
---|
666 | (letrec ((tree-remove |
---|
667 | (lambda (t p-kill) |
---|
668 | |
---|
669 | |
---|
670 | ; (fprintf (current-error-port) "kd-tree-remove: t = ~A~%" t) |
---|
671 | ; (fprintf (current-error-port) "kd-tree-remove: p-kill = ~A~%" p-kill) |
---|
672 | |
---|
673 | (cases kd-tree t |
---|
674 | (KdLeaf (ii pp vv axis) |
---|
675 | |
---|
676 | ; (fprintf (current-error-port) "kd-tree-remove (KdLeaf): vv = ~A~%" vv) |
---|
677 | |
---|
678 | (if vv |
---|
679 | (let ((ipvs (filter-map |
---|
680 | (lambda (i p v) (and (equal? p p-kill) (list i p v))) |
---|
681 | (reverse (cis:elements ii)) pp vv))) |
---|
682 | |
---|
683 | ; (fprintf (current-error-port) "kd-tree-remove (KdLeaf): ipvs = ~A~%" ipvs) |
---|
684 | ; (fprintf (current-error-port) "kd-tree-remove (KdLeaf): pp = ~A~%" pp) |
---|
685 | ; (fprintf (current-error-port) "kd-tree-remove (KdLeaf): vv = ~A~%" vv) |
---|
686 | |
---|
687 | (and (pair? ipvs) |
---|
688 | (let ((ii1 (fold (lambda (i ax) (cis:remove i ax)) |
---|
689 | ii (map car ipvs))) |
---|
690 | (pp1 (fold (lambda (x ax) (remove (lambda (p) (equal? p x)) ax)) |
---|
691 | pp (map cadr ipvs))) |
---|
692 | (vv1 (fold (lambda (x ax) |
---|
693 | (remove (lambda (p) (equal? p x)) ax)) |
---|
694 | vv (map caddr ipvs))) |
---|
695 | ) |
---|
696 | ; (fprintf (current-error-port) "kd-tree-remove (KdLeaf): ii1 = ~A~%" ii1) |
---|
697 | |
---|
698 | (KdLeaf ii1 pp1 vv1 axis)) |
---|
699 | )) |
---|
700 | |
---|
701 | (let ((ips (filter-map (lambda (i p) (and (equal? p p-kill) (list i p))) |
---|
702 | (reverse (cis:elements ii)) pp))) |
---|
703 | (and (pair? ips) |
---|
704 | (let ((ii1 (fold (lambda (i ax) (cis:remove i ax)) |
---|
705 | ii (map car ips))) |
---|
706 | (pp1 (fold (lambda (x ax) (remove (lambda (p) (equal? p x)) ax)) |
---|
707 | pp (map cadr ips))) |
---|
708 | ) |
---|
709 | (KdLeaf ii1 pp1 vv axis)) |
---|
710 | )) |
---|
711 | )) |
---|
712 | |
---|
713 | (KdNode (l p i v r axis ci) |
---|
714 | |
---|
715 | ; (fprintf (current-error-port) "kd-tree-remove (KdNode): p = ~A~%" p) |
---|
716 | ; (fprintf (current-error-port) "kd-tree-remove (KdNode): p-kill = ~A~%" p-kill) |
---|
717 | |
---|
718 | (if (equal? p p-kill) |
---|
719 | |
---|
720 | (let ((offset (if (kd-tree-empty? l) (+ i 1) (kd-tree-min-index l)))) |
---|
721 | (if v |
---|
722 | (let ((pts1 (append (kd-tree->list* l) (kd-tree->list* r)))) |
---|
723 | (list->kd-tree* 0 (length pts1) pts1 axis offset: offset)) |
---|
724 | (let ((pts1 (append (kd-tree->list l) (kd-tree->list r)))) |
---|
725 | (list->kd-tree 0 (length pts1) pts1 axis offset: offset)) |
---|
726 | )) |
---|
727 | |
---|
728 | (if (< (coord axis p-kill) |
---|
729 | (coord axis p)) |
---|
730 | |
---|
731 | (let* ((l1 (tree-remove l p-kill)) |
---|
732 | (r1 (or (and l1 r) (tree-remove r p-kill))) |
---|
733 | (ll (or l1 l)) |
---|
734 | (rr (or r1 r))) |
---|
735 | |
---|
736 | (and (or l1 r1) |
---|
737 | (let ((min1 (cond ((and (kd-tree-empty? ll) |
---|
738 | (kd-tree-empty? rr)) |
---|
739 | i) |
---|
740 | ((kd-tree-empty? ll) |
---|
741 | (kd-tree-min-index rr)) |
---|
742 | (else |
---|
743 | (kd-tree-min-index ll)))) |
---|
744 | (max1 (cond ((and (kd-tree-empty? ll) |
---|
745 | (kd-tree-empty? rr)) |
---|
746 | i) |
---|
747 | ((kd-tree-empty? rr) |
---|
748 | (kd-tree-max-index ll)) |
---|
749 | (else |
---|
750 | (kd-tree-max-index rr)))) |
---|
751 | ) |
---|
752 | |
---|
753 | (KdNode ll p i v rr axis |
---|
754 | (cis:add i (cis:interval min1 max1))) |
---|
755 | ))) |
---|
756 | |
---|
757 | (let* ((r1 (tree-remove r p-kill)) |
---|
758 | (l1 (or (and r1 l) (tree-remove l p-kill))) |
---|
759 | (ll (or l1 l)) |
---|
760 | (rr (or r1 r)) |
---|
761 | ) |
---|
762 | |
---|
763 | (and (or r1 l1) |
---|
764 | (let ((min1 (cond ((and (kd-tree-empty? ll) |
---|
765 | (kd-tree-empty? rr)) |
---|
766 | i) |
---|
767 | ((kd-tree-empty? ll) |
---|
768 | (kd-tree-min-index rr)) |
---|
769 | (else |
---|
770 | (kd-tree-min-index ll)))) |
---|
771 | (max1 (cond ((and (kd-tree-empty? ll) |
---|
772 | (kd-tree-empty? rr)) |
---|
773 | i) |
---|
774 | ((kd-tree-empty? rr) |
---|
775 | (kd-tree-max-index ll)) |
---|
776 | (else |
---|
777 | (kd-tree-max-index rr)))) |
---|
778 | ) |
---|
779 | (KdNode ll p i v rr axis |
---|
780 | (cis:add i (cis:interval min1 max1))) |
---|
781 | ))) |
---|
782 | ))) |
---|
783 | )) |
---|
784 | )) |
---|
785 | tree-remove))) |
---|
786 | |
---|
787 | |
---|
788 | ;; Checks whether the K-D tree property holds for a given tree. |
---|
789 | ;; |
---|
790 | ;; Specifically, it tests that all points in the left subtree lie to |
---|
791 | ;; the left of the plane, p is on the plane, and all points in the |
---|
792 | ;; right subtree lie to the right. |
---|
793 | |
---|
794 | (define=> (make-kd-tree-is-valid? <Point>) |
---|
795 | (lambda (t) |
---|
796 | (cases kd-tree t |
---|
797 | (KdLeaf (ii pp vv axis) #t) |
---|
798 | |
---|
799 | (KdNode (l p i v r axis ci) |
---|
800 | (let ((x (coord axis p))) |
---|
801 | (and (every (lambda (y) (<= (coord axis y) x )) |
---|
802 | (kd-tree->list l)) |
---|
803 | (every (lambda (y) (>= (coord axis y) x)) |
---|
804 | (kd-tree->list r))))) |
---|
805 | ))) |
---|
806 | |
---|
807 | |
---|
808 | ;; Checks whether the K-D tree property holds for the given tree and |
---|
809 | ;; all subtrees. |
---|
810 | |
---|
811 | (define (make-kd-tree-all-subtrees-are-valid? kd-tree-is-valid?) |
---|
812 | (lambda (t) (every kd-tree-is-valid? (kd-tree-subtrees t)))) |
---|
813 | |
---|
814 | |
---|
815 | (define=> (make-kd-tree-slice <Point>) |
---|
816 | (lambda (x-axis x1 x2 t) |
---|
817 | (let recur ((t t) (pts '())) |
---|
818 | (cases kd-tree t |
---|
819 | |
---|
820 | (KdLeaf (ii pp vv axis) |
---|
821 | (append (filter (lambda (p) |
---|
822 | (and (<= x1 (coord x-axis p)) |
---|
823 | (<= (coord x-axis p) x2))) |
---|
824 | pp) |
---|
825 | pts)) |
---|
826 | |
---|
827 | |
---|
828 | (KdNode (l p i v r axis ci) |
---|
829 | (if (= axis x-axis) |
---|
830 | |
---|
831 | (cond ((and (<= x1 (coord axis p)) |
---|
832 | (<= (coord axis p) x2)) |
---|
833 | (recur l (cons p (recur r pts)))) |
---|
834 | |
---|
835 | ((< (coord axis p) x1) |
---|
836 | (recur r pts)) |
---|
837 | |
---|
838 | ((< x2 (coord axis p)) |
---|
839 | (recur l pts))) |
---|
840 | |
---|
841 | (if (and (<= x1 (coord x-axis p)) |
---|
842 | (<= (coord x-axis p) x2)) |
---|
843 | (recur l (cons p (recur r pts))) |
---|
844 | (recur l (recur r pts))) |
---|
845 | )) |
---|
846 | )) |
---|
847 | )) |
---|
848 | |
---|
849 | |
---|
850 | (define=> (make-kd-tree-slice* <Point>) |
---|
851 | (lambda (x-axis x1 x2 t) |
---|
852 | (let recur ((t t) (pts '())) |
---|
853 | (cases kd-tree t |
---|
854 | (KdLeaf (ii pp vv axis) |
---|
855 | (append |
---|
856 | (if vv |
---|
857 | (filter-map (lambda (i v p) |
---|
858 | (and (<= x1 (coord x-axis p)) |
---|
859 | (<= (coord x-axis p) x2) |
---|
860 | (list (list i v) p))) |
---|
861 | (reverse (cis:elements ii)) vv pp) |
---|
862 | (filter-map (lambda (i p) |
---|
863 | (and (<= x1 (coord x-axis p)) |
---|
864 | (<= (coord x-axis p) x2) |
---|
865 | (list i p))) |
---|
866 | (reverse (cis:elements ii)) pp)) |
---|
867 | pts)) |
---|
868 | |
---|
869 | (KdNode (l p i v r axis ci) |
---|
870 | (if (= axis x-axis) |
---|
871 | |
---|
872 | (cond ((and (<= x1 (coord axis p)) |
---|
873 | (<= (coord axis p) x2)) |
---|
874 | (recur l (cons (if v (list (list i v) p) (list i p)) (recur r pts)))) |
---|
875 | |
---|
876 | ((< (coord axis p) x1) |
---|
877 | (recur r pts)) |
---|
878 | |
---|
879 | ((< x2 (coord axis p)) |
---|
880 | (recur l pts))) |
---|
881 | |
---|
882 | (if (and (<= x1 (coord x-axis p)) |
---|
883 | (<= (coord x-axis p) x2)) |
---|
884 | (recur l (cons (if v (list (list i v) p) (list i p)) (recur r pts))) |
---|
885 | (recur l (recur r pts))) |
---|
886 | )) |
---|
887 | )) |
---|
888 | )) |
---|
889 | |
---|
890 | (define (default-<KdTree> point-class) |
---|
891 | (let* ((list->kd-tree/depth |
---|
892 | (make-list->kd-tree/depth point-class)) |
---|
893 | (kd-tree-remove |
---|
894 | ((make-kd-tree-remove point-class) |
---|
895 | (list->kd-tree/depth identity #f) |
---|
896 | (list->kd-tree/depth cadr (lambda (i v) (cadar v))))) |
---|
897 | (kd-tree-nearest-neighbor |
---|
898 | (make-kd-tree-nearest-neighbor point-class))) |
---|
899 | |
---|
900 | (make-<KdTree> |
---|
901 | (lambda (points #!key |
---|
902 | (leaf-factor 10) |
---|
903 | (make-point identity) |
---|
904 | (make-value #f) |
---|
905 | (offset 0) |
---|
906 | ) |
---|
907 | ((list->kd-tree/depth make-point make-value) |
---|
908 | 0 (length points) points 0 |
---|
909 | leaf-factor: leaf-factor offset: offset)) |
---|
910 | |
---|
911 | (make-kd-tree-nearest-neighbor point-class) |
---|
912 | (make-kd-tree-nearest-neighbor* point-class) |
---|
913 | (make-kd-tree-near-neighbors point-class) |
---|
914 | (make-kd-tree-near-neighbors* point-class) |
---|
915 | ((make-kd-tree-k-nearest-neighbors point-class) |
---|
916 | kd-tree-remove kd-tree-nearest-neighbor) |
---|
917 | kd-tree-remove |
---|
918 | (make-kd-tree-slice point-class) |
---|
919 | (make-kd-tree-slice* point-class) |
---|
920 | (make-kd-tree-is-valid? point-class) |
---|
921 | (make-kd-tree-all-subtrees-are-valid? |
---|
922 | (make-kd-tree-is-valid? point-class)) |
---|
923 | ))) |
---|
924 | |
---|
925 | (define KdTree3d |
---|
926 | (default-<KdTree> Point3d)) |
---|
927 | |
---|
928 | (define KdTree2d |
---|
929 | (default-<KdTree> Point2d)) |
---|
930 | |
---|
931 | |
---|
932 | |
---|
933 | ) |
---|
934 | |
---|