Skip to content

k-d tree

之前在做 Golearn 時,為了加速 knn 的搜尋速度,因此實做了 k-d tree,當時找了很久,總是找不到一篇完整的文章講解 如何 將 kdtree 應用在 knn 中 ,因此決定寫這篇來紀錄下我的作法。

什麼是 k-d tree ?

k-d tree 是 k-dimensional tree 的縮寫,也就是 k 維樹,簡單來說,就是樹的每一層都是以 不同的維度標準 做分割,而不是單純用某個維度來建樹。

單維度的樹

通常看到的樹,都是只用其中一個維度的資料來建樹,舉例來說:

有六筆資料如下,每筆資料都只有一個維度

[1, 3, 5, 41, 67, 87]

直接建樹的話,可能會出現像這樣的樹

多維度的樹(k-d tree)

通常用在機器學習的資料不會只有一個維度,舉例來說:

一樣有六筆資料,每筆資料都是二維的

[[2, 4],
 [3, 7],
 [6, 5],
 [1, 6],
 [8, 4],
 [5, 3]]

如果每個維度的資料都一樣重要的話,建樹時就必須考慮所有的維度,因此在建 k-d tree 時,就會 輪流用不同的維度 作為切割的標準

k-d tree 分析

將上面的資料畫到二維座標圖上來分析 其中紅色的點是上面給的資料,淺藍色的線是第一個分界,紫色的線是第二個分界

藉由 x 方向以及 y 方向的線,將所有的點都區分開來 如果現在有一個新的點,我想要知道 它跟誰最近 的話,就可以利用 k-d tree 快速找到屬於那個點的區域

建構 k-d tree

建樹

建樹要考慮兩個部份,分別是 要用哪個維度分割 以及 要選哪個點當分割節點

  • 維度的部份,我是直接照著順序從第一個開始每層就換成下一個,值得注意的是,其實 不一定 要照順序來,倒著做或是隨機選一個都是可以的
  • 選擇節點的部份,我是將當前區塊的所有點依照選定維度的值,由小到大排列,接著選取 中間點 作為分割節點
    • 其餘的若選定維度的值小於等於分割節點,則放入左子樹
    • 若大於分割節點,則放入右子樹

我在查 k-d tree 相關的資料時,有一些資料說要把所有的點都下放到葉節點上,但我搞不懂這樣如何實做,因此選擇把分割節點掛在上面的作法

尋找最近點

如果現在有一個新的點,我想要知道離它最近的點時,並不是直接往下看到哪個點停下來就是最近的,而是要 一層一層尋找,雖說是一層一層找,但並不會搜尋所有的點

具體操作如下:

  1. 往下比較直到找到最底層的節點
  2. 計算出當前點的距離
  3. 若當前點距離小於最小距離,則更新最小距離,並搜尋該節點另一子節點下所有節點
  4. 移動到父節點
  5. 重複執行 2, 3, 4 步直到到達根節點

knn 使用 k-d tree

所謂的 knn 指的是 k nearest neighbors,顧名思義,knn 就是要在一堆資料中找出離某資料最近的 k 筆

最簡單也最耗時的作法就是全部都計算出距離後,再排列後拿出最近的 k 個,這個作法稱作 linear

實際上,也可以利用 k-d tree 以在不用檢查所有資料的情況下,快速找出最近的k筆資料,但問題是上述的 k-d tree 只能找到最近的一個點,因此需要做一點修改

heap

在我的作法中會使用到 max-heap,所謂的 max-heap 就是會將丟入的資料存下來,然後可以輸出其中最大值的一種資料結構

前 k 個最小值搜尋

我的作法是使用一個容量為 k 的 heap,在搜尋的過程中

  • 如果 heap 還未滿,則加入當前的節點,且搜尋當前節點的另一個子樹
  • 如果 heap 已滿,則利用 heap 中的當前最大值來執行原本的搜尋,並隨時更新 heap 中的值

具體操作如下:

  1. 往下比較直到找到最底層的節點
  2. 計算出當前點的距離
  3. 若heap未滿,則加入當前節點,且搜尋該節點另一子節點下所有節點
  4. 若heap已滿,且當前點距離小於heap中最大距離,則更新heap,並搜尋該節點另一子節點下所有節點
  5. 移動到父節點
  6. 重複執行2,3,4,5步直到到達根節點

結束後,heap中的k筆資料就是距離最近的k筆了

上述所說的”更新heap”指的是將當前heap中的最大值丟掉,並將新的值放入heap中

後記

參考連結