Kd-tree的用法__演算法基礎

Kd-Tree演算法原理和開源實現代碼

 

本文介紹一種用於高維空間中的快速最近鄰和近似最近鄰尋找技術——Kd-Tree（Kd樹）。Kd-Tree，即K-dimensional tree，是一種高維索引樹形資料結構，常用於在大規模的高維資料空間進行最近鄰尋找(Nearest Neighbor)和近似最近鄰尋找(Approximate Nearest Neighbor)，例如映像檢索和識別中的高維映像特徵向量的K近鄰尋找與匹配。本文首先介紹Kd-Tree的基本原理，然後對基於BBF的近似尋找方法進行介紹，最後給出一些參考文獻和開源實現代碼。

 

一、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）若它的左子樹不為空白，則左子樹上所有結點的值均小於它的根結點的值；

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

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

例如，圖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更近的點（如茶色圈圈內的兩個紅色點）

 

回溯：