Play With KDTree

^{— source code in this page lies here}

最近工作中实现了基于KDTree的3D空间划分，用于实时判断Mesh和Ray的相交，我现在依然惊讶于这个算法的优美和高效。

还是先从工作中逃离一会，试试用Clojure实现一下2D-KDtree，然后在quil框架来些与工作无关的东西看看。

先生成3000个点吧：

(def pts
  (repeatedly 3000 #(vec2 (q/random (q/width))
                          (q/random (q/height)))))

调用quil的point函数画出来看看

 (doall (map #(apply q/point %) pts))

好了， 现在我们可以将这些点通过KDTree算法来进行二维分割了:

(defn divide [pts rect depth]
  (when (>= (count pts) 1)
    (let [top-left (:p rect)
          axis ([:x :y] (mod depth 2))
          sort-pts (vec (sort-by #(axis %) pts))
          median (quot (count pts) 2)
          translate #(g/translate % (m/- top-left))
          left (map translate (subvec sort-pts 0 median))
          right (map translate (subvec sort-pts (inc median)))
          div-rects (divide-rect
                     rect axis
                     (get-in sort-pts [median axis]))]
      {:pts pts
       :rect rect
       :depth depth
       :left (divide left (div-rects 0) (inc depth))
       :right (divide right (div-rects 1) (inc depth))})))

这其实大致于Wikipedia描述的一致，第一个参数pts是当前待分割的点集，第二个参数rect则是该点集包围矩形，第三参数depth则是分割深度。针对当前点集将其以x轴或y轴分割成两部分子点集并得到他们各自的包围矩形，然后递归调用divide直至最后一个点，这样就得到了kdtree数据。

当中有个特殊的处理：

translate #(g/translate % (m/- top-left))
left (map translate (subvec sort-pts 0 median))
right (map translate (subvec sort-pts (inc median)))

这里会将分割出来的左子集和右子集的点映射到以其父包围矩形的Left-Top为坐标系原点来表示，换言之，每次分割我们都可以理解成得到一个子坐标系，这样做的目的主要是为了方便后面对每一层做缩放变换。

(def kdtree (divide pts (rect (q/width) (q/height)) 0)))

好了，这样就生成了该3000个点集的kdtree，我们可以检测看看深度为9的右子树的数据：

(assoc (get-in kdtree (repeat 9 :right)) :left "..." :right "...")

=> {:pts
     ([42.2945556640625 6.13824462890625]
      [50.94134521484375 40.2008056640625]
      [56.525146484375 41.60552978515625]
      [57.4437255859375 13.4390869140625]),
     :rect
     {:p [35.60064697265625 0.0],
      :size [36.08294677734375 52.22308349609375]},
      :depth 9,
      :left "...",
      :right "..."}

可以看到， 划分到第9层时只剩下4个点了。

当中每一层的包围矩形是我很感兴趣的，那么我们将该kdtree所有的包围矩形画出来看看， 代码很简单：

(defn draw-kdtree [tree]
  (when tree
    (q/push-matrix)  
    (let [rect (:rect tree)
          size (:size rect)]
      (q/translate (:p rect))
      (q/rect 0 0 (:x size) (:y size))
      (draw-kdtree (:left tree))
      (draw-kdtree (:right tree)))
    (q/pop-matrix)))

(draw-kdtree kdtree)

颜色和缩放 #

KDTree本身就呈现出很有意思的秩序感，我可以再加些颜色和缩放看看，假设想要针对每一层包围矩形实施中心缩放操作，并使其内部子包围矩形也会被相应缩放，貌似比较麻烦，但因为上面在生成KDTree数据的时候已经将每一层划分的包围矩形映射到子坐标系了，所以只需要针对每一层子坐标系做中心缩放即可。其次，要填充包围矩形很简单，直接挑选喜欢的颜色调用Fill接口就可以绘制填充色了， 这里我们只填充叶子节点的包围矩形，这两部操作的实现只需要将上面的draw-kdtree拓展为:

(defn draw-kdtree [tree]                      
  (when tree                                  
    (q/push-matrix)                           
    (let [rect (:rect tree)                   
          size (:size rect)]                  
      (q/translate (:p rect))
      (transform-center                       
       (m/div size 2)                         
       (q/scale (q/random 0.8 1.2)))
      (if (is-leaf tree)                      
        (do (q/fill (rand-color))
            (q/rect 0 0 (:x size) (:y size)))
        (do (draw-kdtree (:left tree))        
            (draw-kdtree (:right tree)))))    
    (q/pop-matrix)))