longchute

about

k-d Tree

07 May 2012

NOTE: I originally posted this on Snipplr.

make-kd-node [median left right]: Creates a node in a kd-tree.

make-kd-tree [k depth points]: Creates a kd-tree of kd-nodes. TODO: Not stack safe. Use loop/recur.

kd-nearest-neighbor [a-point kd-tree]: Returns the nearest neighboring point to a given point, using a kd-tree.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
(ns kd-tree
    (:use [clojure.zip]))

(defn ^{
    :doc "Creates a node in a kd-tree."
} make-kd-node [median left right] [median left right])

(defn ^{
    :doc "Creates a kd-tree of kd-nodes. **TODO**: Not stack safe. Use loop/recur."
} make-kd-tree [k depth points]
    (if (empty? points)
        nil
        (let [axis          (mod depth k)
              sorted-points (sort-by (fn [a-point] (nth a-point axis)) > points)
              pivot         (/ (count points) 2)
              median        (nth sorted-points pivot)
              left          (take (dec pivot) sorted-points)
              right         (drop pivot sorted-points)
              next-depth    (inc depth)]
                (make-kd-node 
                    median
                    (make-kd-tree k next-depth left) 
                    (make-kd-tree k next-depth right)))))

(defn ^{
    :doc "Returns the nearest neighboring point to a given point, using a kd-tree."
} kd-nearest-neighbor [a-point kd-tree]
    (let [kd-tree-zipper (seq-zip kd-tree)]
        (-> kd-tree-zipper down)))
                    
(kd-nearest-neighbor [3 3 3] 
    (make-kd-tree 3 0 [[0 1 2]
                       [1 2 0]
                       [2 0 1]
                       [3 4 5]
                       [5 3 4]
                       [4 5 3]]))