本文介紹一種用于高維空間中的快速最近鄰和近似最近鄰查找技術——Kd-Tree（Kd樹）。

Kd-Tree，即K-dimensional tree，是一種高維索引樹形數據結構，常用于在大規模的高維數據空間進行最近鄰查找(Nearest Neighbor)和近似最近鄰查找(Approximate Nearest Neighbor)，例如圖像檢索和識別中的高維圖像特征向量的K近鄰查找與匹配。本文首先介紹Kd-Tree的基本原理，然后對基于BBF的近似查找方法進行介紹，最后給出一些參考文獻和開源實現代碼。

1. Kd-tree

Kd-Tree，即K-dimensional tree，是一棵二叉樹，樹中存儲的是一些K維數據。在一個K維數據集合上構建一棵Kd-Tree代表了對該K維數據集合構成的K維空間的一個劃分，即樹中的每個結點就對應了一個K維的超矩形區域（Hyperrectangle）。

在介紹Kd-tree的相關算法前，我們先回顧一下二叉查找樹（Binary Search Tree）的相關概念和算法。

二叉查找樹（Binary Search Tree，BST），是具有如下性質的二叉樹（來自wiki）：

若它的左子樹不為空，則左子樹上所有結點的值均小于它的根結點的值；

若它的右子樹不為空，則右子樹上所有結點的值均大于它的根結點的值；

它的左、右子樹也分別為二叉排序樹；

例如，圖1中是一棵二叉查找樹，其滿足BST的性質。

圖1 二叉查找樹（來源：Wiki）

給定一個1維數據集合，怎樣構建一棵BST樹呢？根據BST的性質就可以創建，即將數據點一個一個插入到BST樹中，插入后的樹仍然是BST樹，即根結點的左子樹中所有結點的值均小于根結點的值，而根結點的右子樹中所有結點的值均大于根結點的值。

將一個1維數據集用一棵BST樹存儲后，當我們想要查詢某個數據是否位于該數據集合中時，只需要將查詢數據與結點值進行比較然后選擇對應的子樹繼續往下查找即可，查找的平均時間復雜度為：O(logN)，最壞的情況下是O(N)。

如果我們要處理的對象集合是一個K維空間中的數據集，那么是否也可以構建一棵類似于1維空間中的二叉查找樹呢？答案是肯定的，只不過推廣到K維空間后，創建二叉樹和查詢二叉樹的算法會有一些相應的變化（后面會介紹到兩者的區別），這就是下面我們要介紹的Kd-tree算法。

怎樣構造一棵Kd-tree？

對于Kd-tree這樣一棵二叉樹，我們首先需要確定怎樣劃分左子樹和右子樹，即一個K維數據是依據什么被劃分到左子樹或右子樹的。

在構造1維BST樹時，一個1維數據根據其與樹的根結點和中間結點進行大小比較的結果來決定是劃分到左子樹還是右子樹，同理，我們也可以按照這樣的方式，將一個K維數據與Kd-tree的根結點和中間結點進行比較，只不過不是對K維數據進行整體的比較，而是選擇某一個維度Di，然后比較兩個K維數在該維度Di上的大小關系，即每次選擇一個維度Di來對K維數據進行劃分，相當于用一個垂直于該維度Di的超平面將K維數據空間一分為二，平面一邊的所有K維數據在Di維度上的值小于平面另一邊的所有K維數據對應維度上的值。也就是說，我們每選擇一個維度進行如上的劃分，就會將K維數據空間劃分為兩個部分，如果我們繼續分別對這兩個子K維空間進行如上的劃分，又會得到新的子空間，對新的子空間又繼續劃分，重復以上過程直到每個子空間都不能再劃分為止。以上就是構造Kd-Tree的過程，上述過程中涉及到兩個重要的問題：

1、每次對子空間的劃分時，怎樣確定在哪個維度上進行劃分。

2、在某個維度上進行劃分時，怎樣確保在這一維度上的劃分得到的兩個子集合的數量盡量相等，即左子樹和右子樹中的結點個數盡量相等。

問題1：每次對子空間的劃分時，怎樣確定在哪個維度上進行劃分？

最簡單的方法就是輪著來，即如果這次選擇了在第i維上進行數據劃分，那下一次就在第j(j≠i)維上進行劃分，例如：j = (i mod k) + 1。想象一下我們切豆腐時，先是豎著切一刀，切成兩半后，再橫著來一刀，就得到了很小的方塊豆腐。

可是“輪著來”的方法是否可以很好地解決問題呢？再次想象一下，我們現在要切的是一根木條，按照“輪著來”的方法先是豎著切一刀，木條一分為二，干凈利落，接下來就是再橫著切一刀，這個時候就有點考驗刀法了，如果木條的直徑（橫截面）較大，還可以下手，如果直徑較小，就沒法往下切了。因此，如果K維數據的分布像上面的豆腐一樣，“輪著來”的切分方法是可以奏效，但是如果K維度上數據的分布像木條一樣，“輪著來”就不好用了。因此，還需要想想其他的切法。

如果一個K維數據集合的分布像木條一樣，那就是說明這K維數據在木條較長方向代表的維度上，這些數據的分布散得比較開，數學上來說，就是這些數據在該維度上的方差（invariance）比較大，換句話說，正因為這些數據在該維度上分散的比較開，我們就更容易在這個維度上將它們劃分開，因此，這就引出了我們選擇維度的另一種方法：最大方差法（max invarince），即每次我們選擇維度進行劃分時，都選擇具有最大方差維度。

問題2：在某個維度上進行劃分時，怎樣確保在這一維度上的劃分得到的兩個子集合的數量盡量相等，即左子樹和右子樹中的結點個數盡量相等？

假設當前我們按照最大方差法選擇了在維度i上進行K維數據集S的劃分，此時我們需要在維度i上將K維數據集合S劃分為兩個子集合A和B，子集合A中的數據在維度i上的值都小于子集合B中。首先考慮最簡單的劃分法，即選擇第一個數作為比較對象（即劃分軸，pivot），S中剩余的其他所有K維數據都跟該pivot在維度i上進行比較，如果小于pivot則劃A集合，大于則劃入B集合。把A集合和B集合分別看做是左子樹和右子樹，那么我們在構造一個二叉樹的時候，當然是希望它是一棵盡量平衡的樹，即左右子樹中的結點個數相差不大。而A集合和B集合中數據的個數顯然跟pivot值有關，因為它們是跟pivot比較后才被劃分到相應的集合中去的。好了，現在的問題就是確定pivot了。給定一個數組，怎樣才能得到兩個子數組，這兩個數組包含的元素個數差不多且其中一個子數組中的元素值都小于另一個子數組呢？方法很簡單，找到數組中的中值（即中位數，median），然后將數組中所有元素與中值進行比較，就可以得到上述兩個子數組。同樣，在維度i上進行劃分時，pivot就選擇該維度i上所有數據的中值，這樣得到的兩個子集合數據個數就基本相同了。

解決了上面兩個重要的問題后，就得到了Kd-Tree的構造算法了。

Kd-Tree的構建算法：

（1） 在K維數據集合中選擇具有最大方差的維度k，然后在該維度上選擇中值m為pivot對該數據集合進行劃分，得到兩個子集合；同時創建一個樹結點node，用于存儲；

（2）對兩個子集合重復（1）步驟的過程，直至所有子集合都不能再劃分為止；如果某個子集合不能再劃分時，則將該子集合中的數據保存到葉子結點（leaf node）。

以上就是創建Kd-Tree的算法。下面給出一個簡單例子。

給定二維數據集合：(2,3), (5,4), (9,6), (4,7), (8,1), (7,2)，利用上述算法構建一棵Kd-tree。左圖是Kd-tree對應二維數據集合的一個空間劃分，右圖是構建的一棵Kd-tree。

圖2 構建的kd-tree

其中圓圈代表了中間結點(k, m)，而紅色矩形代表了葉子結點。

Kd-Tree與一維二叉查找樹之間的區別：

二叉查找樹：數據存放在樹中的每個結點（根結點、中間結點、葉子結點）中；

Kd-Tree：數據只存放在葉子結點，而根結點和中間結點存放一些空間劃分信息（例如劃分維度、劃分值）；

構建好一棵Kd-Tree后，下面給出利用Kd-Tree進行最近鄰查找的算法：

（1）將查詢數據Q從根結點開始，按照Q與各個結點的比較結果向下訪問Kd-Tree，直至達到葉子結點。

其中Q與結點的比較指的是將Q對應于結點中的k維度上的值與m進行比較，若Q(k) < m，則訪問左子樹，否則訪問右子樹。達到葉子結點時，計算Q與葉子結點上保存的數據之間的距離，記錄下最小距離對應的數據點，記為當前“最近鄰點”Pcur和最小距離Dcur。

（2）進行回溯（Backtracking）操作，該操作是為了找到離Q更近的“最近鄰點”。即判斷未被訪問過的分支里是否還有離Q更近的點，它們之間的距離小于Dcur。

如果Q與其父結點下的未被訪問過的分支之間的距離小于Dcur，則認為該分支中存在離P更近的數據，進入該結點，進行（1）步驟一樣的查找過程，如果找到更近的數據點，則更新為當前的“最近鄰點”Pcur，并更新Dcur。

如果Q與其父結點下的未被訪問過的分支之間的距離大于Dcur，則說明該分支內不存在與Q更近的點。

回溯的判斷過程是從下往上進行的，直到回溯到根結點時已經不存在與P更近的分支為止。

怎樣判斷未被訪問過的樹分支Branch里是否還有離Q更近的點？

從幾何空間上來看，就是判斷以Q為中心center和以Dcur為半徑Radius的超球面（Hypersphere）與樹分支Branch代表的超矩形（Hyperrectangle）之間是否相交。

在實現中，我們可以有兩種方式來求Q與樹分支Branch之間的距離。第一種是在構造樹的過程中，就記錄下每個子樹中包含的所有數據在該子樹對應的維度k上的邊界參數[min, max]；第二種是在構造樹的過程中，記錄下每個子樹所在的分割維度k和分割值m，（k, m），Q與子樹的距離則為|Q(k) - m|。

以上就是Kd-tree的構造過程和基于Kd-Tree的最近鄰查找過程。

下面用一個簡單的例子來演示基于Kd-Tree的最近鄰查找的過程。

數據點集合：(2,3), (4,7), (5,4), (9,6), (8,1), (7,2) 。

已建好的Kd-Tree：

圖3 構建的kd-tree

其中，左圖中紅色點表示數據集合中的所有點。

查詢點：(8, 3) （在左圖中用茶色菱形點表示）

第一次查詢：

圖4 第一次查詢的kd-tree

當前最近鄰點：(9, 6) ， 最近鄰距離：sqrt(10)，

且在未被選擇的樹分支中存在于Q更近的點（如茶色圈圈內的兩個紅色點）

回溯：

圖5 回溯kd-tree

當前最近鄰點：(8, 1)和(7, 2) ， 最近鄰距離：sqrt(2)

最后，查詢點(8, 3)的近似最近鄰點為(8, 1)和(7, 2) 。

2. Kd-tree with BBF

上一節介紹的Kd-tree在維度較小時（例如：K≤30），算法的查找效率很高，然而當Kd-tree用于對高維數據（例如：K≥100）進行索引和查找時，就面臨著維數災難（curse of dimension）問題，查找效率會隨著維度的增加而迅速下降。通常，實際應用中，我們常常處理的數據都具有高維的特點，例如在圖像檢索和識別中，每張圖像通常用一個幾百維的向量來表示，每個特征點的局部特征用一個高維向量來表征（例如：128維的SIFT特征）。因此，為了能夠讓Kd-tree滿足對高維數據的索引，Jeffrey S. Beis和David G. Lowe提出了一種改進算法——Kd-tree with BBF（Best Bin First），該算法能夠實現近似K近鄰的快速搜索，在保證一定查找精度的前提下使得查找速度較快。

在介紹BBF算法前，我們先來看一下原始Kd-tree是為什么在低維空間中有效而到了高維空間后查找效率就會下降。在原始kd-tree的最近鄰查找算法中（第一節中介紹的算法），為了能夠找到查詢點Q在數據集合中的最近鄰點，有一個重要的操作步驟：回溯，該步驟是在未被訪問過的且與Q的超球面相交的子樹分支中查找可能存在的最近鄰點。隨著維度K的增大，與Q的超球面相交的超矩形（子樹分支所在的區域）就會增加，這就意味著需要回溯判斷的樹分支就會更多，從而算法的查找效率便會下降很大。

一個很自然的思路是：既然kd-tree算法在高維空間中是由于過多的回溯次數導致算法查找效率下降的話，我們就可以限制查找時進行回溯的次數上限，從而避免查找效率下降。這樣做有兩個問題需要解決：1）最大回溯次數怎么確定？2）怎樣保證在最大回溯次數內找到的最近鄰比較接近真實最近鄰，即查找準確度不能下降太大。

問題1）：最大回溯次數怎么確定？

最大回溯次數一般人為設定，通常根據在數據集上的實驗結果進行調整。

問題2）：怎樣保證在最大回溯次數內找到的最近鄰比較接近真實最近鄰，即查找準確度不能下降太大？

限制回溯次數后，如果我們還是按照原來的回溯方法挨個地進行訪問的話，那很顯然最后的查找結果的精度就很大程度上取決于數據的分布和回溯次數了。挨個訪問的方法的問題在于認為每個待回溯的樹分支中存在最近鄰的概率是一樣的，所以它對所有的待回溯樹分支一視同仁。實際上，在這些待回溯樹分支中，有些樹分支存在最近鄰的可能性比其他樹分支要高，因為樹分支離Q點之間的距離或相交程度是不一樣的，離Q更近的樹分支存在Q的最近鄰的可能性更高。因此，我們需要區別對待每個待回溯的樹分支，即采用某種優先級順序來訪問這些待回溯樹分支，使得在有限的回溯次數中找到Q的最近鄰的可能性很高。我們要介紹的BBF算法正是基于這樣的解決思路，下面我們介紹BBF查找算法。

基于BBF的Kd-Tree近似最近鄰查找算法

已知：

Q：查詢數據；? KT：已建好的Kd-Tree；

1. 查找Q的當前最近鄰點P

1）從KT的根結點開始，將Q與中間結點node(k,m)進行比較，根據比較結果選擇某個樹分支Branch（或稱為Bin）；并將未被選擇的另一個樹分支（Unexplored Branch）所在的樹中位置和它跟Q之間的距離一起保存到一個優先級隊列中Queue；

2）按照步驟1）的過程，對樹分支Branch進行如上比較和選擇，直至訪問到葉子結點，然后計算Q與葉子結點中保存的數據之間的距離，并記錄下最小距離D以及對應的數據P。

注：

A、Q與中間結點node(k,m)的比較過程：如果Q(k) > m則選擇右子樹，否則選擇左子樹。

B、優先級隊列：按照距離從小到大的順序排列。

C、葉子結點：每個葉子結點中保存的數據的個數可能是一個或多個。

2. 基于BBF的回溯

已知：最大回溯次數BTmax

1）如果當前回溯的次數小于BTmax，且Queue不為空，則進行如下操作：

從Queue中取出最小距離對應的Branch，然后按照1.1步驟訪問該Branch直至達到葉子結點；計算Q與葉子結點中各個數據間距離，如果有比D更小的值，則將該值賦給D，該數據則被認為是Q的當前近似最近鄰點；

2）重復1）步驟，直到回溯次數大于BTmax或Queue為空時，查找結束，此時得到的數據P和距離D就是Q的近似最近鄰點和它們之間的距離。

下面用一個簡單的例子來演示基于Kd-Tree+BBF的近似最近鄰查找的過程。

數據點集合：(2,3), (4,7), (5,4), (9,6), (8,1), (7,2) 。

已建好的Kd-Tree：

圖6 構建的kd-tree

基于BBF的查找的過程：

查詢點Q：? (5.5, 5)

第一遍查詢：

圖7 第一次查詢的kd-tree

當前最近鄰點：(9, 6) ， 最近鄰距離：sqrt(13.25)，

同時將未被選擇的樹分支的位置和與Q的距離記錄到優先級隊列中。

BBF回溯：

從優先級隊列里選擇距離Q最近的未被選擇樹分支進行回溯。

圖8 利用BBF方法回溯kd-tree

當前最近鄰點：(4, 7) ， 最近鄰距離：sqrt(6.25)

繼續從優先級隊列里選擇距離Q最近的未被選擇樹分支進行回溯。

圖9 利用BBF方法回溯kd-tree

當前最近鄰點：(5, 4) ， 最近鄰距離：sqrt(1.25)

最后，查詢點(5.5, 5)的近似最近鄰點為(5, 4) 。

原創[Python數據科學](https://mp.weixin.qq.com/s/lnXB6_73JiOrYL2cFefYiw)