1. 原理篇

我們用大白話講講KD-Tree是怎么一回事。

1.1 線性查找

假設數組A為[0, 6, 3, 8, 7, 4, 11]，有一個元素x，我們要找到數組A中距離x最近的元素，應該如何實現呢？比較直接的想法是用數組A中的每一個元素與x作差，差的絕對值最小的那個元素就是我們要找的元素。假設x = 2，那么用數組A中的所有元素與x作差得到[-2, 4, 1, 6, 5, 2, 9]，其中絕對值最小的是1，對應的元素是數組A中的3，所以3就是我們的查找結果。

1.2 二分查找

如果我們有大量的元素要在數組A中進行查找，那么1.1的方式就顯得不是那么高效了，如果數組A的長度為N，那么每次查找都要進行N次操作，即算法復雜度為O(N)。

我們把數組A進行升序排列，得到[0, 3, 4, 6, 7, 8, 11]；

令x = 2，數組中間的元素是6，2小于6，所以2只可能存在于6的左邊，我們只需要在數組[0, 3, 4]中繼續查找；

左邊的數組中間的元素是3，2小于3，所以2只可能存在于3的左邊，即數組[0]；

由于數組[0]無法再分割，查找結束；

x需要跟我們最終找到的0，以及倒數第二步找到的3進行比較，發現2離3更近，所以查找結果為3。
這種查找方法就是二分查找，其算法復雜度為O(Log2(N))。

1.3 BST

除了數組之外，有沒有更直觀的數據結構可以實現1.2的二分查找呢？答案就是二分查找樹，全稱Binary Search Tree，簡稱BST。把數組A建立成一個BST，結構如下圖所示。我們只需要訪問根節點，進行值比較來確定下一節點，如此循環往復直到訪問到葉子節點為止。

1.4 多維數組

現在我們把問題加點難度，假設數組B為[[6, 2], [6, 3], [3, 5], [5, 0], [1, 2], [4, 9], [8, 1]]，有一個元素x，我們要找到數組B中距離x最近的元素，應該如何實現呢？比較直接的想法是用數組B中的每一個元素與x求距離，距離最小的那個元素就是我們要找的元素。假設x = [1, 1]，那么用數組A中的所有元素與x求距離得到[5.0, 5.4, 4.5, 4.1, 1.0, 8.5, 7.0]，其中距離最小的是1，對應的元素是數組B中的[1, 2]，所以[1, 2]就是我們的查找結果。

1.5 再次陷入困境

如果我們有大量的元素要在數組B中進行查找，那么1.4的方式就又顯得不是那么高效了，如果數組B的長度為N，那么每次查找都要進行N次操作，即算法復雜度為O(N)。

1.6 什么是KD-Tree

這時候已經沒辦法用BST，不過我們可以對BST做一些改變來適應多維數組的情況。當當當當~，這時候該KD-Tree出場了。廢話不多說，先上圖：

1.7 如何建立KD-Tree

您可能會問，剛在那張圖的KD Tree又是如何建立的呢？ 很簡單，只需要5步：
1. 建立根節點；

2. 選取方差最大的特征作為分割特征；

3. 選擇該特征的中位數作為分割點；

4. 將數據集中該特征小于中位數的傳遞給根節點的左兒子，大于中位數的傳遞給根節點的右兒子；

5. 遞歸執行步驟2-4，直到所有數據都被建立到KD Tree的節點上為止。

不難看出，KD Tree的建立步驟跟BST是非常相似的，可以認為BST是KD Tree在一維數據上的特例。KD Tree的算法復雜度介于O(Log2(N))和O(N)之間。

1.8 特征選取

您可能還會問，為什么方差最大的適合作為特征呢？ 因為方差大，數據相對“分散”，選取該特征來對數據集進行分割，數據散得更“開”一些。

1.9 分割點選擇

您可能又要問，為什么選擇中位數作為分割點呢？ 因為借鑒了BST，選取中位數，讓左子樹和右子樹的數據數量一致，便于二分查找。

1.10 利用KD-Tree查找元素

KD Tree建好之后，接下來就要利用KD Tree對元素進行查找了。查找的方式在BST的基礎上又增加了一些難度，如下：
1. 從根節點開始，根據目標在分割特征中是否小于或大于當前節點，向左或向右移動。

2. 一旦算法到達葉節點，它就將節點點保存為“當前最佳”。

3. 回溯，即從葉節點再返回到根節點

4. 如果當前節點比當前最佳節點更接近，那么它就成為當前最好的。

5. 如果目標距離當前節點的父節點所在的將數據集分割為兩份的超平面的距離更接近，說明當前節點的兄弟節點所在的子樹有可能包含更近的點。因此需要對這個兄弟節點遞歸執行1-4步。

1.11 超平面

所以什么是超平面呢，聽起來讓人一臉懵逼。
以[0, 2, 0], [1, 4, 3], [2, 6, 1]的舉例：
1. 如果用第二維特征作為分割特征，那么從三個數據點中的對應特征取出2, 4, 6，中位數是4；

2. 所以[1, 4, 3]作為分割點，將[0, 2, 0]劃分到左邊，[2, 6, 1]劃分到右邊；

3. 從立體幾何的角度考慮，三維空間得用一個二維的平面才能把空間一分為二，這個平面可以用y = 4來表示；

4. 點[0, 2, 0]到超平面y = 4的距離就是 sqrt((2 - 4) ^ 2) = 2；

5. 點[2, 6, 1]到超平面y = 4的距離就是 sqrt((6 - 4) ^ 2) = 2。

2. 實現篇

本人用全宇宙最簡單的編程語言——Python實現了KD-Tree算法，沒有依賴任何第三方庫，便于學習和使用。簡單說明一下實現過程，更詳細的注釋請參考本人github上的代碼。

2.1 創建Node類

初始化，存儲父節點、左節點、右節點、特征及分割點。

class Node(object):def __init__(self):self.father = Noneself.left = Noneself.right = Noneself.feature = Noneself.split = None

2.2 獲取Node的各個屬性

def __str__(self):return "feature: %s, split: %s" % (str(self.feature), str(self.split))

2.3 獲取Node的兄弟節點

@property def brother(self):if self.father is None:ret = Noneelse:if self.father.left is self:ret = self.father.rightelse:ret = self.father.leftreturn ret

2.4 創建KDTree類

初始化，存儲根節點。

class KDTree(object):def __init__(self):self.root = Node()

2.5 獲取KDTree屬性

便于我們查看KD Tree的節點值，各個節點之間的關系。

def __str__(self):ret = []i = 0que = [(self.root, -1)]while que:nd, idx_father = que.pop(0)ret.append("%d -> %d: %s" % (idx_father, i, str(nd)))if nd.left is not None:que.append((nd.left, i))if nd.right is not None:que.append((nd.right, i))i += 1return "\n".join(ret)

2.6 獲取數組中位數的下標

def _get_median_idx(self, X, idxs, feature):n = len(idxs)k = n // 2col = map(lambda i: (i, X[i][feature]), idxs)sorted_idxs = map(lambda x: x[0], sorted(col, key=lambda x: x[1]))median_idx = list(sorted_idxs)[k]return median_idx

2.7 計算特征的方差

注意這里用到了方差公式，D(X) = E(X^2)-[E(X)]^2

def _get_variance(self, X, idxs, feature):n = len(idxs)col_sum = col_sum_sqr = 0for idx in idxs:xi = X[idx][feature]col_sum += xicol_sum_sqr += xi ** 2return col_sum_sqr / n - (col_sum / n) ** 2

2.8 選擇特征

取方差最大的的特征作為分割點特征。

def _choose_feature(self, X, idxs):m = len(X[0])variances = map(lambda j: (j, self._get_variance(X, idxs, j)), range(m))return max(variances, key=lambda x: x[1])[0]

2.9 分割特征

把大于、小于中位數的元素分別放到兩個列表中。

def _split_feature(self, X, idxs, feature, median_idx):idxs_split = [[], []]split_val = X[median_idx][feature]for idx in idxs:if idx == median_idx:continuexi = X[idx][feature]if xi < split_val:idxs_split[0].append(idx)else:idxs_split[1].append(idx)return idxs_split

2.10 建立KDTree

使用廣度優先搜索的方式建立KD Tree，注意要對X進行歸一化。

def build_tree(self, X, y):X_scale = min_max_scale(X)nd = self.rootidxs = range(len(X))que = [(nd, idxs)]while que:nd, idxs = que.pop(0)n = len(idxs)if n == 1:nd.split = (X[idxs[0]], y[idxs[0]])continuefeature = self._choose_feature(X_scale, idxs)median_idx = self._get_median_idx(X, idxs, feature)idxs_left, idxs_right = self._split_feature(X, idxs, feature, median_idx)nd.feature = featurend.split = (X[median_idx], y[median_idx])if idxs_left != []:nd.left = Node()nd.left.father = ndque.append((nd.left, idxs_left))if idxs_right != []:nd.right = Node()nd.right.father = ndque.append((nd.right, idxs_right))

2.11 搜索輔助函數

比較目標元素與當前結點的當前feature，訪問對應的子節點。反復執行上述過程，直到到達葉子節點。

def _search(self, Xi, nd):while nd.left or nd.right:if nd.left is None:nd = nd.rightelif nd.right is None:nd = nd.leftelse:if Xi[nd.feature] < nd.split[0][nd.feature]:nd = nd.leftelse:nd = nd.rightreturn nd

2.12 歐氏距離

計算目標元素與某個節點的歐氏距離，注意get_euclidean_distance這個函數沒有進行開根號的操作，所以求出來的是歐氏距離的平方。

def _get_eu_dist(self, Xi, nd):X0 = nd.split[0]return get_euclidean_distance(Xi, X0)

2.13 超平面距離

計算目標元素與某個節點所在超平面的歐氏距離，為了跟2.11保持一致，要加上平方。

def _get_hyper_plane_dist(self, Xi, nd):j = nd.featureX0 = nd.split[0]return (Xi[j] - X0[j]) ** 2

2.14 搜索函數

搜索KD Tree中與目標元素距離最近的節點，使用廣度優先搜索來實現。

def nearest_neighbour_search(self, Xi):dist_best = float("inf")nd_best = self._search(Xi, self.root)que = [(self.root, nd_best)]while que:nd_root, nd_cur = que.pop(0)while 1:dist = self._get_eu_dist(Xi, nd_cur)if dist < dist_best:dist_best = distnd_best = nd_curif nd_cur is not nd_root:nd_bro = nd_cur.brotherif nd_bro is not None:dist_hyper = self._get_hyper_plane_dist(Xi, nd_cur.father)if dist > dist_hyper:_nd_best = self._search(Xi, nd_bro)que.append((nd_bro, _nd_best))nd_cur = nd_cur.fatherelse:breakreturn nd_best

3 效果評估

3.1 線性查找

用“笨”辦法查找距離最近的元素。

def exhausted_search(X, Xi):dist_best = float('inf')row_best = Nonefor row in X:dist = get_euclidean_distance(Xi, row)if dist < dist_best:dist_best = distrow_best = rowreturn row_best

3.2 main函數

主函數分為如下幾個部分：

1. 隨機生成數據集，即測試用例

2. 建立KD-Tree

3. 執行“笨”辦法查找

4. 比較“笨”辦法和KD-Tree的查找結果

def main():def main():print("Testing KD Tree...")test_times = 100run_time_1 = run_time_2 = 0for _ in range(test_times):low = 0high = 100n_rows = 1000n_cols = 2X = gen_data(low, high, n_rows, n_cols)y = gen_data(low, high, n_rows)Xi = gen_data(low, high, n_cols) tree = KDTree()tree.build_tree(X, y)start = time()nd = tree.nearest_neighbour_search(Xi)run_time_1 += time() - startret1 = get_euclidean_distance(Xi, nd.split[0])start = time()row = exhausted_search(X, Xi)run_time_2 += time() - startret2 = get_euclidean_distance(Xi, row)assert ret1 == ret2, "target:%s\nrestult1:%s\nrestult2:%s\ntree:\n%s" \% (str(Xi), str(nd), str(row), str(tree)) print("%d tests passed!" % test_times) print("KD Tree Search %.2f s" % run_time_1) print("Exhausted search %.2f s" % run_time_2)

https://github.com/tushushu/imylu/tree/master/imylu/utils

總結

以上是生活随笔為你收集整理的KD Tree的原理及Python实现的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。