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