144 lines
5.4 KiB
Common Lisp
144 lines
5.4 KiB
Common Lisp
(in-package cl-user)
|
|
|
|
(defvar *random-target*
|
|
(list (float (/ (random 1000) 100))
|
|
(float (/ (random 1000) 100))
|
|
(float (/ (random 1000) 100))))
|
|
|
|
(defun distance (n target &key (semi nil))
|
|
"distance node to target:
|
|
returns squared euclidean distance, or squared semi distance if option set"
|
|
(if semi
|
|
(expt (- (nth (first n) (second n))
|
|
(nth (first n) target))
|
|
2)
|
|
(reduce #'+
|
|
(mapcar (lambda (x y) (* (- x y) (- x y)))
|
|
(second n)
|
|
target))))
|
|
|
|
(defun target< (n target)
|
|
"returns true if target is to its left in axis dim"
|
|
(< (nth (first n) target)
|
|
(nth (first n) (second n))))
|
|
|
|
(defun next-node (n target &key (opposite nil))
|
|
"return the next child when nn searching, return opposing child if option
|
|
oppose set"
|
|
(if (or (and (target< n target) (not opposite))
|
|
(and (not (target< n target)) opposite))
|
|
(third n)
|
|
(fourth n)))
|
|
|
|
(defun make-kdtree (axis data dims)
|
|
"a kdtree is a binary tree where nodes are:
|
|
terminal: (axis data-point), or
|
|
branch: (axis split-point (left-kdtree) (right-kdtree))"
|
|
(cond ((null data) nil)
|
|
((eql (length data) 1) ; singleton?
|
|
(list axis (first data))) ; terminal node
|
|
;; else branch node:
|
|
;; #pts=odd splits list into 2 even parts with sp in middle
|
|
;; #pts=even splits list into 2 uneven parts with shorter length first (but never nil)
|
|
(t
|
|
(let ((sd (sort (copy-list data)
|
|
#'<
|
|
:key (lambda (x) (nth axis x)))) ; sort the axis ordinates
|
|
(sp (truncate (length data) 2)) ; get mid pt
|
|
(nxta (mod (1+ axis) dims)))
|
|
(list axis
|
|
(nth sp sd)
|
|
(make-kdtree nxta
|
|
(subseq sd 0 sp)
|
|
dims)
|
|
(make-kdtree nxta
|
|
(subseq sd (1+ sp))
|
|
dims))))))
|
|
|
|
(defun visit-kdtree (kdt &key (node-function nil))
|
|
"depth first visit all nodes in kdtree and optionally apply a function to
|
|
each node visited"
|
|
(when kdt
|
|
(when node-function (funcall node-function kdt))
|
|
(visit-kdtree (third kdt) :node-function node-function)
|
|
(visit-kdtree (fourth kdt) :node-function node-function)))
|
|
|
|
(defun count-nodes (kdt)
|
|
"count of the terminal nodes"
|
|
(cond ((null kdt) 0)
|
|
((= (length kdt) 2) 1)
|
|
(t (+ 1
|
|
(count-nodes (third kdt))
|
|
(count-nodes (fourth kdt))))))
|
|
|
|
(defvar *hits*)
|
|
|
|
(defun nn-kdtree (kdt node-stack target)
|
|
"nearest neighbour search"
|
|
(when kdt
|
|
;; stage 1 - find the 'closest' terminal node using insertion logic
|
|
(let ((best (do ((node kdt (next-node node target)))
|
|
((not (next-node node target))
|
|
(incf *hits*)
|
|
node) ; return first best est.
|
|
(push node node-stack)
|
|
(incf *hits*))))
|
|
;; stage 2 - unwind the path, at each node if node is closer then make it best
|
|
(do ((node (pop node-stack) (pop node-stack)))
|
|
((null node) best) ; return nearest pt
|
|
;; iteration: update best if node is closer
|
|
(when (< (distance node target) (distance best target))
|
|
(setf best node))
|
|
;; venture down opposing side if split point is inside HS
|
|
(let ((opposing-best
|
|
(if (< (distance node target :semi t) (distance best target))
|
|
(nn-kdtree (next-node node target :opposite t)
|
|
(list)
|
|
target)
|
|
nil))) ; otherwise ignore this subtree
|
|
(when (and opposing-best
|
|
(< (distance opposing-best target) (distance best target)))
|
|
(setf best opposing-best)))))))
|
|
|
|
(defun process (data target
|
|
&key (render nil)
|
|
&aux (dims (length target)))
|
|
"process one set of data & optionally display tree"
|
|
(let* ((*hits* 0)
|
|
(kdt (make-kdtree 0 data dims))
|
|
(nn (nn-kdtree kdt (list) target)))
|
|
(when render
|
|
(visit-kdtree kdt
|
|
:node-function
|
|
(lambda (n)
|
|
(format t
|
|
"~A node: axis:~A point: ~A target:~A semi-distance-sqd:~A euclidean-distance-sqd:~A~%"
|
|
(if (not (next-node n target))
|
|
"TERMINAL"
|
|
"BRANCH")
|
|
(first n)
|
|
(second n)
|
|
target
|
|
(distance n target :semi t)
|
|
(distance n target)))))
|
|
(format t "~%NN to ~A is ~A, distance ~A [tree has ~A nodes, ~A were visited.]~%"
|
|
target
|
|
(second nn)
|
|
(sqrt (distance nn target))
|
|
(count-nodes kdt)
|
|
*hits*)))
|
|
|
|
(defun main ()
|
|
;; TASK 1 - nn search small set of 2D points
|
|
(process '((2 3) (5 4) (9 6) (4 7) (8 1) (7 2))
|
|
'(9 2)
|
|
:render t)
|
|
|
|
;; TASK 2 - nn search 1000 coordinate points in 3D space
|
|
(process (let ((ll (list)))
|
|
(dotimes (i 10)
|
|
(dotimes (j 10)
|
|
(dotimes (k 10)
|
|
(push (list i j k) ll))))
|
|
ll)
|
|
*random-target*))
|