68 lines
2.3 KiB
Racket
68 lines
2.3 KiB
Racket
#lang racket
|
|
|
|
; A tree consists of a point, a left and a right subtree.
|
|
(struct tree (p l r) #:transparent)
|
|
; If the node is in depth d, then the points in l has
|
|
; the (d mod k)'th coordinate less than the same coordinate in p.
|
|
|
|
(define (kdtree d k ps)
|
|
(cond [(empty? ps) #f] ; #f represents an empty subtree
|
|
[else (define-values (p l r) (split-points ps (modulo d k)))
|
|
(tree p (kdtree (+ d 1) k l) (kdtree (+ d 1) k r))]))
|
|
|
|
(define (split-points ps d)
|
|
(define (ref p) (vector-ref p d))
|
|
(define sorted-ps (sort ps < #:key ref))
|
|
(define mid (quotient (+ (length ps)) 2))
|
|
(define median (ref (list-ref sorted-ps mid)))
|
|
(define-values (l r) (partition(λ(x)(< (ref x) median))sorted-ps))
|
|
(values (first r) l (rest r)))
|
|
|
|
; The bounding box of a subtree:
|
|
(struct bb (mins maxs) #:transparent)
|
|
|
|
(define (infinite-bb k)
|
|
(bb (make-vector k -inf.0) (make-vector k +inf.0)))
|
|
|
|
(define/match (copy-bb h)
|
|
[((bb mins maxs))
|
|
(bb (vector-copy mins) (vector-copy maxs))])
|
|
|
|
(define (dist v w) (for/sum ([x v] [y w]) (sqr (- x y))))
|
|
(define (intersects? g r hr) (<= (dist (closest-in-hr g hr) g) r))
|
|
(define (closest-in-hr g hr)
|
|
(for/vector ([gi g] [mini (bb-mins hr)] [maxi (bb-maxs hr)])
|
|
(cond [(<= gi mini) mini]
|
|
[(< mini gi maxi) gi]
|
|
[else maxi])))
|
|
|
|
(define (split-bb hr d x)
|
|
(define left (copy-bb hr))
|
|
(define right (copy-bb hr))
|
|
(vector-set! (bb-maxs left) d x)
|
|
(vector-set! (bb-mins right) d x)
|
|
(values left right))
|
|
|
|
(define visits 0) ; for statistics only
|
|
(define (visit) (set! visits (+ visits 1)))
|
|
(define (reset-visits) (set! visits 0))
|
|
(define (regret-visit) (set! visits (- visits 1)))
|
|
|
|
(define (nearest-neighbor g t k)
|
|
(define (nearer? p q) (< (dist p g) (dist q g)))
|
|
(define (nearest p q) (if (nearer? p q) p q))
|
|
(define (nn d t bb) (visit)
|
|
(define (ref p) (vector-ref p (modulo d k)))
|
|
(match t
|
|
[#f (regret-visit) #(+inf.0 +inf.0 +inf.0)]
|
|
[(tree p l r)
|
|
(define-values (lbb rbb) (split-bb bb (modulo d k) (ref p)))
|
|
(define-values (near near-bb far far-bb)
|
|
(if (< (ref g) (ref p))
|
|
(values l lbb r rbb)
|
|
(values r rbb l lbb)))
|
|
(define n (nearest p (nn (+ d 1) near near-bb)))
|
|
(if (intersects? g (dist n g) far-bb)
|
|
(nearest n (nn (+ d 1) far far-bb))
|
|
n)]))
|
|
(nn 0 t (infinite-bb k)))
|