KD-Tree 演算法筆記

簡介

KD 樹又稱 K 維樹 (K-dimensional tree)，是一種可以對 K 維資料進行劃分的資料結構，可以看成二元搜尋樹的一種延伸，不斷的對空間中的維度做劃分，利用搜尋樹剪枝的特性縮短時間複雜度，主要應用在多維空間搜尋，例如最近鄰居搜索。

背景

以前在學資料結構的時候，學過使用二元搜尋樹將資料用大小區分來建樹，每個點代表著一筆資料，這個方法在資料只有一個維度的時候行的通，但當資料超過一個維度的時候就遇到困難，該怎麼對二維以上的陣列進行劃分。

想法

對於多維陣列，我們多了一個以上的維度，因此在劃分時沒有一個劃分依據，此時我們可以將所有資料統一使用其中一個維度進行劃分，這個動作在二維資料內相當於將空間劃分為兩個部分，得到兩個新的子空間，如果我們繼續對這兩個子空間進行上述劃分，又會得到新的子空間，重複以上過程直到每個子空間都不能再劃分為止。

相當於不斷的用超平面將 K 維空間切分，每個節點對應一個 K 維超矩形區域。

以上就是建立 KD 樹的過程，上述過程中涉及到兩個重要的問題：

1. 每次選擇其中一個維度進行劃分時，應該選擇哪個維度進行劃分？
2. 在某個維度上的進行劃分時，應該選擇哪個節點進行劃分？

劃分方法

術語

• cell：KD 樹數據結構為把空間遞歸地劃分兩個不相交的超矩形（類似於超平面、超立方體的概念），每個超矩形所佔的空間即為一個 cell。
• bounding rectangle：包含數據集中所有點的最小超矩形。其範圍小於等於該節點對應 cell 的大小。
• aspect ratio：cells 中最長邊與最短邊的長度比，比值越大，數目 cell 越細長，反之，越肥胖。
• spread：數據集在某維度上最大值減最小值。即該維度上 bounding rectangle 對應邊的長度。

準則

以下列出 ANN 中提供的幾種劃分規則：

• standard kd-tree splitting rule（標準 KD 樹分割規則）：
  • 基本思想：使創建出來的樹盡可能的保持平衡，樹越平衡代表著分割得越平均，搜尋的時間也就越少。
  • 維度選擇有兩種方法（二擇一）：
    1. 選擇 spread 最大的維度，即 bounding rectangle 中最長邊對應的維度。即數據集在該維度上，最大值減最小值的結果為所有維度中最大的維度。
    2. 選擇變異數最大的維度，即數據在該維度上，數據分散的比較開的維度，
  • 劃分點選擇：選擇數據集在該維度上的中位數做為劃分點，讓左子樹和右子樹的元素數量盡量保持一致，且其中一個子節點的元素值都小於另一個子節點。
  • 優點：劃分到左右子樹中數據數量相同（或相差 1），為平衡二元樹。
  • 缺點：
    1. 維度選擇方法一，可能出現細長的cell，容易增加回溯時遞迴搜尋的範圍。
    2. 維度選擇方法二，計算最大變異數的時間複雜度要高於計算最大 spread 的時間複雜度。

---

• midpoint splitting rule（中點分割規則）：
  • 基本思想：讓分割出的 cells 盡可能的減少產生細長的 cell，使劃分出的兩個 cells 都是肥胖的。
  • 維度選擇：cell 中最長邊對應的維度。
  • 劃分點選擇：cell 中最長邊對應的中心。
  • 優點：找中心比找中值的時間複雜度更低。
  • 缺點：導致數據集劃分不均勻，當數據大量聚集時，存在很多無用的空間劃分（即劃分時，所有數據都在一個區域，另一個區域沒有數據）。

---

• sliding midpoint rule（滑動中點規則）：
  • 基本思想：midpoint splitting rule 的改進版，規則大致相同，其認為兩種劃分為好的劃分：
    1. 平衡的劃分：使得劃分出的兩個子 cell 都是肥胖的（midpoint splitting rule中的思想）
    2. 使得較肥胖的 cell 中含有較少的數據點。（對無效劃分的處理）
  • 維度選擇：選擇超矩形中最長邊對應的維度。（同 midpoint splitting rule）
  • 劃分點的選擇：在劃分時，當劃分結果沒有包含任何數據集，出現無效劃分時，通過滑動分割點來消除無效劃分，若該維度的中點值為 a：
    • 若所有數據在該維度上都大於 a，則以數據集中最小值來劃分空間，使左子樹區域中有 1 個點，右子樹有 n-1 個點。
    • 若所有數據在該維度上都小於 a，則以數據集中最大值來劃分空間，使左子樹區域中有 n-1 個點，右子樹有 1 個點。
  • 優點：通過滑動劃分點，可以消除無效的劃分。其可能存在細長的 cell，但其始終伴隨著一個該維度上較寬胖的 cell。

---

• fair-split rule（公平分割規則）：
  • standard kd-tree splitting rule 與 midpoint splitting rule 的折衷方法
  • 基本思想：設定一個合理的 aspect ratio 值為 a。在劃分空間時盡量滿足子 cell 的 aspect ratio 小於 a 且盡量使數據均勻的分佈在兩個子 cell 中。
  • 給定一個 cell：
    • 根據 cell 各邊的長度，選出能夠使得劃分出兩個子 cell 能滿足 aspect ratio 小於 a 的所有維度。
    • 在上述選出的維度中，選擇 bounding rectangle 中最長邊對應的維度。
    • 在滿足 aspect ratio 小於 a 的前提下，在該維度上選擇一個使得數據點盡可能分佈均勻的劃分點來對 cell 進行最終的劃分。（standard kd-tree splitting rule 的思想 - 取中位數）
  • 優點：相比於 midpoint splitting rule，對於無效的劃分，有更好的劃分結果。然而當數據高度聚集時，仍有可能出現與 midpoint splitting rule 相同的結果。

---

• sliding fair-split rule（滑動公平分割規則）：結合 fair-split rule 與 sliding midpoint rule 的思想
  • 基本思想：上述 fair-split rule 劃分不能避免：由於必須滿足 aspect ratio 小於 a，使得極端情況下，使數據點盡可能的均勻分佈的劃分仍然為一個無效劃分的情況。此時可以通過滑動劃分點至該維度上最大（或最小）的值處（即 sliding 的思想）來進行劃分（此時無需滿足 aspect ratio 小於 a 的條件），使得每個劃分都至少含有一個數據點。

無論哪種劃分規則，都不會影響 KD 樹搜索結果的正確性，只會影響其樹的形狀和深度，從而影響搜索性能。總而言之，希望得到的 KD 樹的結構盡量滿足兩個特點：

1. 盡可能的平衡
2. 盡可能的讓 cell 肥胖（使搜索過程中與超矩形相交的個數盡量少，以減小到韃子葉的次數）

資料儲存方法

將數據儲存在所有節點

• 葉節點、中間節點：
  • 維度
  • 點
  • 左子節點
  • 右子節點

將數據僅儲存在葉節點

• 中間節點：
  • 維度
  • 劃分軸
  • 左子節點
  • 右子節點
• 葉節點：
  • 點

建樹

這裡我們建立 KD 樹的方法使用：

• standard kd-tree splitting rule（標準 KD 樹分割規則）
• 選擇最大變異數做為維度選擇方法
• 將數據儲存在所有節點

決定上述兩個方法後，我們 KD 樹的建立過程如下：

1. 在 K 維資料集合中選擇具有最大變異數的維度 K，然後在該維度上選擇中位數 m 做為分割點對該資料集合進行劃分，得到兩個子集合，同時建立一個樹節點用於儲存。
2. 對兩個子集合重複 (1) 步驟的過程，直到所有子集合都不能在劃分為止。

範例

假設有一個二維陣列：[[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]]

KD 樹演算法就是要確定下圖中這些分割空間的分割線（多維空間即為分割平面，一般為超平面）。

1. 分別計算 x, y 方向上資料的變異數，得出 x 方向上的變異數最大
2. 根據 x 軸方向的值 (2, 4, 5, 7, 8, 9)，排序選出中位數為 5 和 7，這裡選擇 7 當作分割點，所以該節點中的資料為 (7, 2)。這樣該節點的分割超平面就是通過 (7, 2) 並垂直於 x 軸的直線 x = 7。
3. 確定左子空間和右子空間。以 x = 7 為分割超平面將整個空間分為兩個部分，x <= 7 的部分為左子空間，包含 3 個節點 (2, 3), (4, 7), (5, 4)，另一部分為右子空間，包含2個節點 (8, 1), (9, 6)。

KD 樹的建立是一個遞迴的過程。繼續對左子空間和右子空間內的資料重複根節點的操作就可以得到下一級子節點 (5, 4) 和 (9, 6) (也就是左右子空間的根節點)，同時將空間和資料進一步細分。如此反覆值到空間中只包含一個資料點。

尋找最近點

KD Tree 建好之後，接下來就可以使用 KD Tree 對元素搜尋最近點，方法如下：

• 查詢目標節點：Q
• 當前最佳節點：P

1. 建立一個空的堆疊 S。
2. 從根節點開始走訪，根據 Q 在分割維度中是否小於或大於當前節點，向左或向右移動，將每個走訪過的節點都存入（Push） S。
3. 一旦不能再走訪，設定 P 為無限大，並開始循環以下：
   1. 將 S 取出（Pop）為當前節點 C。從概念上來看，S 會不斷取出走訪過的節點，意義即為從葉節點返回父節點，稱為回溯。
   2. 計算 C 到 Q 的距離，如果 C 比 P 更接近 Q，那麼更新 P 為該節點。
   3. 檢查在 C 分割面的另一邊是否有比 P 距離更近的節點。從概念上來說，以 Q 為中心，以 P 為半徑劃一個超球面，看這個超球面是否穿過了分割平面。因為平面都是座標軸對應的，所以只需要簡單比較 Q 和當前節點分割面的距離是否比 P 距離更小。
      1. 如果超球面穿越分割面，那麼分割面的另外一側可能還有最近點，所以需要對 KD 樹的另一邊的分支在執行一次走訪，將每個走訪過的節點都存入（Push） S。
      2. 如果超球面沒有穿過分割面，則分割面的另外一邊的整個分支會被剪掉，稱為剪枝，並繼續搜尋其它節點，即回溯。
4. 當 S 為空，也就是算法最後回溯到根節點的時候，搜尋完成，返回結果。

1. 通常，算法使用距離平方來做比較，而不使用更耗時的平方根。可以通過維持當前最佳的平方距離來節省計算量。
2. 計算距離步驟可以在檢查分割面後面，優點是在判斷完分割面沒有相交的節點後，可以直接判斷跳過計算該距離的步驟，因為只要沒有與分割面相交，距離就不會更近。這裡因為解釋順序方便將距離計算在前面，實際這樣會在某些情況額外計算一些無意義的距離。

剪枝概念

KD 演算法的核心技巧在於剪枝，在已經搜索到 B 時，發現其到 B 的距離，要比到 A 的右子樹的平面距離還更短，所以整個 A 的右子樹都被剪枝，一下子剪去了一半的點。

這個算法也可以通過簡單的修改做多種延伸。比如，可用於計算 K 個最近臨點，這個時候需要保存 K 個當前最佳而不是一個。分支能夠剪掉的條件是：K 個點都找到，並且分支中沒有比這 K 個最佳更近的點。

也可以轉換為近似算法加快運行。例如，近似最近點搜尋可以通過指定檢查節點的上限來達成，也可以基於時鐘(在硬體實現更合適)終止搜尋過程。

範例一

如圖所示，點 Q(2.1, 3.1) 為要查詢的點。

1. 從根節點開始走訪，根據分割維度的數據判斷應該進入左或右子樹
   1. 進入 (7, 2)，分割軸為 x，(2.1 < 7)，進入左子樹
   2. 進入 (5, 4)，分割軸為 y，(3.1 < 4)，進入左子樹
   3. 進入 (2, 3)，分割軸為 x，(2.1 > 2)，進入右子樹
   4. 右子樹為空，因此停止走訪，搜尋路徑依序為 (7, 2) -> (5, 4) -> (2, 3)
2. 根據搜尋路徑尋找最近鄰點：
   1. 取出 (2, 3)，分割軸為 x
      1. 計算到 Q 的距離得出 0.14，更新 P = 0.14
      2. 計算 Q 到 x = 2 的距離，|2.1 - 2| = 0.1
      3. 因為 0.1 < 0.14，與超球面相交，進入 (2, 3) 的另一子樹
         1. 另一子樹為空，因此停止走訪，搜尋路徑依序為 (7, 2) -> (5, 4)
   2. 取出 (5, 4)，分割軸為 y
      1. 計算到 Q 的距離得出 3.03，3.03 > P，不用更新 P
      2. 計算 Q 到 y = 4 的距離，|3.1 - 4| = 0.9
      3. 因為 0.9 > 0.14，沒有與超球面相交，剪掉 (5, 4) 的另一子樹
   3. 取出 (7, 2)，分割軸為 x
      1. 計算到 Q 的距離得出 5.02，5.02 > P，不用更新 P
      2. 計算 Q 到 x = 7 的距離，|2.1 - 7| = 4.9
      3. 因為 4.9 > 0.14，沒有與超球面相交，剪掉 (7, 2) 的另一子樹
3. 至此搜尋路徑中的節點已經全部回溯完畢，返回：
   • 最近鄰點 P(2, 3)
   • 最近距離為 0.1414

範例二

如下圖所示，點 Q(2, 4.5) 為要查詢的點。

1. 從根節點開始走訪，根據分割維度的數據判斷應該進入左或右子樹
   1. 進入 (7, 2)，分割軸為 x，(2 < 7)，進入左子樹
   2. 進入 (5, 4)，分割軸為 y，(4.5 > 4)，進入右子樹
   3. 進入 (4, 7)，分割軸為 x，(2 < 4)，進入左子樹
   4. 左子樹為空，因此停止走訪，搜尋路徑依序為 (7, 2) -> (5, 4) -> (4, 7)
2. 根據搜尋路徑尋找最近鄰點：
   1. 取出 (4, 7)，分割軸為 x
      1. 計算到 Q 的距離得出 3.20，更新 P = 3.20
      2. 計算 Q 到 x = 4 的距離，|2 - 4| = 2
      3. 因為 2 < 3.20，與超球面相交，進入 (4, 7) 的另一子樹
      4. 另一子樹為空，因此停止走訪
   2. 取出 (5, 4)，分割軸為 y
      1. 計算到 Q 的距離得出 3.04，3.04 < P，P = 3.04
      2. 計算 Q 到 y = 4 的距離，|4.5 - 4| = 0.5
      3. 因為 0.5 < 3.04，與超球面相交，進入 (5, 4) 的另一子樹
         1. 進入 (2, 3)，分割軸為 x，(2 <= 2)，進入左子樹
         2. 左子樹為空，因此停止走訪，搜尋路徑依序為 (7, 2) -> (2, 3)
   3. 取出 (2, 3)，分割軸為 x
      1. 計算到 Q 的距離得出 1.5，1.5 < P，P = 1.5
      2. 計算 Q 到 x = 2 的距離，|2 - 2| = 0
      3. 因為 0 < 1.5，與超球面相交，進入 (2, 3) 的另一子樹
         1. 另一子樹為空，因此停止走訪，搜尋路徑依序為 (7, 2)
   4. 取出 (7, 2)，分割軸為 x
      1. 計算到 Q 的距離得出 5.59，5.59 > P，不用更新 P
      2. 計算 Q 到 x = 7 的距離，|2 - 7| = 5
      3. 因為 5 > 1.5，沒有與超球面相交，剪掉 (7, 2) 的另一子樹
3. 至此搜尋路徑中的節點已經全部回溯完畢，返回：
   • 最近鄰點 P(2, 3)
   • 最近距離為 1.5

添加節點

和 BST 樹添加節點的方式一樣。首先，從根節點開始判斷從哪個維度開始分割的，再依據該維度的元素確定要進入左節點還是右節點，直到找到葉節點，將待添加的節點依據該分割維度元素的大小進入左節點或右節點中。

按照這種方式添加節點可能會導致樹失去平衡，從而降低樹的性能。樹性能降低的比例取決於樹之前的空間分布，以及添加的節點樹和樹原大小的關係。如果樹變得很不平衡，就需要做樹的平衡，從而恢復依賴於樹平衡的查詢性能，例如最近臨查詢。

刪除節點

從已經建立好的樹之中刪除節點，且不破壞限制條件，最簡單的方法是將待刪除的節點及其子樹做成集合，並重新建立子樹。

另外一個方法是為待刪除節點找一個替代點。首先，找到包含待刪除點的節點R，如果 R 是葉節點，則不需要替換，直接刪除R。如果是其他情況，從以 R 為根的子樹中找到一個替代點，設為 p，交換 R 和 p，再刪除 p。

找到一個可替換點的方法：假設節點 R 通過 x 軸來區分，並且 R 有一個右節點，找到這個右節點及其子樹中 x 值最小的點， 即可做為替換點。反之，找到左節點及其子樹中 x 值最大的點，即可做為替換點。

平衡樹

KD 樹的平衡需要非常小心，因為 KD 樹通過多個維度來排序，所以樹旋轉這樣的技術不能用來做平衡，原因是這個技術可能破壞 KD 樹的限制條件。

KD 樹有幾種變體，包括：divided k-d tree, pseudo k-d tree, K-D-B-tree。這些變體裡面有許多是自適應 KD 樹。

範圍搜索

範圍搜索指的是使用範圍參數來做檢索，例如，如果一個 KD 樹儲存的是收入和年齡的樹值，那麼一個範圍搜索可能是：查找樹中年齡在 20 到 50，收入在 50000 到 80000 的節點。 應為 KD 樹在樹的每一層對域的範圍做了分割，所以可以高效執行範圍查詢。

高維數據會降低性能

在隨機分布的數據點上，搜尋最近點時間複雜度是 O(log(N))，在高維度空間，維數災難會導致算法需要訪問遠多於低維空間的分支。在實踐中，如果節點的數量僅略高於維度的數量時，該算法僅比對所有點搜索的線性搜索略好。

在高維空間，KD樹是不適合做高效的近鄰搜尋。通常原則是，如果維度是 K，樹據點樹是 N，需要滿足 N >> 2^k。否則，當 KD 樹用在高維度數據上，搜尋時絕大多數節點需要做評估，所以性能不一定比窮舉搜索好，應該替換為一個近似的近鄰搜尋。

參考文獻

• ANN (Approximate Nearest Neighbor) Programming Manual
• It’s okay to be skinny, if your friends are fat
• k-d tree
• Point kd-tree
• Kd Tree
• Kd Tree 演算法詳解
• K - 近鄰算法 KNN 學習筆記
• KD 樹 / K-D 樹
• KD-TREE 演算法原理
• KD Tree 的原理及 Python 實現
• KD-Tree 原理詳解
• 【量化課堂】kd 樹算法之詳細篇