k-d tree

(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]]))