機(jī)器學(xué)習(xí)常用的KD-Tree深度講解

1、從線段樹到KD樹在講KD樹之前，我們先來(lái)了解一下線段樹的概念。線段樹在機(jī)器學(xué)習(xí)領(lǐng)域 當(dāng)中不太常見(jiàn)，作為高性能維護(hù)的數(shù)據(jù)結(jié)構(gòu)，經(jīng)常出現(xiàn)在各種算法比賽當(dāng)中。 線段樹的本質(zhì)是一棵維護(hù)一段區(qū)間的平衡二叉樹。比如下圖就是一個(gè)經(jīng)典的線 段樹：2 3X 1從下圖當(dāng)中我們不難看出來(lái)，這棵線段樹維護(hù)的是個(gè)區(qū)間內(nèi)的最大值。 比如樹根是8,維護(hù)的是整個(gè)區(qū)間的最大值，每一個(gè)中間節(jié)點(diǎn)的值都是以它為 樹根的子樹中所有元素的最大值。通過(guò)線段樹，我們可以在的時(shí)間內(nèi)計(jì)算出某一個(gè)連續(xù)區(qū)間的最大值。比如 我們來(lái)看下圖：b當(dāng)我們要求被框起來(lái)的區(qū)間中的最大值，我們只需要找到能夠覆蓋這個(gè)區(qū) 間的中間節(jié)點(diǎn)就行。我們可以發(fā)現(xiàn)被紅框框起來(lái)的兩

2、個(gè)節(jié)點(diǎn)的子樹剛好覆蓋這 個(gè)區(qū)間，于是整個(gè)區(qū)間的最大值，就是這兩個(gè)元素的最大值。這樣，我們就把 一個(gè)需要查找的問(wèn)題降低成了，不但如此，我們也可以做到復(fù)雜度內(nèi)的更新， 也就是說(shuō)我們不但可以快速查詢，還可以更新線段當(dāng)中的元素。當(dāng)然線段樹的 應(yīng)用非常廣泛，也有許多種變體，這里我們不過(guò)多深入，感興趣的同學(xué)可以期 待一下周三的算法與數(shù)據(jù)結(jié)構(gòu)專題，在之后的文章當(dāng)中會(huì)為大家分享線段樹的 相關(guān)內(nèi)容。在這里，我們只需要有一個(gè)大概的印象，線段樹究竟完成的是什么 樣的事情即可。線段樹維護(hù)的是一個(gè)線段，也就是區(qū)間內(nèi)的元素，也就是說(shuō)維護(hù)的是一個(gè)一維的序列。如果我們將數(shù)據(jù)的維度擴(kuò)充一下，擴(kuò)充到多維呢？是 的，你沒(méi)有猜錯(cuò)，從

3、某種程度上來(lái)說(shuō)，我們可以把KD-Tree看成是線段樹拓展到多維空間當(dāng)中的情況。KD-Tree 定義我們來(lái)看一下KD-Tree的具體定義，這里的 K指的是K維空間，D自然就 是dimension，也就是維度，也就是說(shuō) KD-Tree就是K維度樹的意思。在我們 構(gòu)建線段樹的時(shí)候，其實(shí)是一個(gè)遞歸的建樹過(guò)程，我們每次把當(dāng)前的線段一分 為二，然后用分成兩半的數(shù)據(jù)分別構(gòu)建左右子樹。我們可以簡(jiǎn)單寫一下偽代碼， 來(lái)更直觀地感受一下：class Node:def _init_(self, value, lchild, rchild):self.value = valueself.lchild = lchilds

4、elf.rchild = rchild def build(arr):n = len( arr):left, right = arr: n /2, arrn /2:Ichild, rchild = self.build(left), self.build(right)return Node(max(lchild.value, rchild.value), Ichild, rchild)我們來(lái)看一個(gè)二維的例子，在一個(gè)二維的平面當(dāng)中分布著若干個(gè)點(diǎn)。x軸。我們對(duì)我們首先選擇一個(gè)維度將這些數(shù)據(jù)一分為二，比如我們選擇 所有數(shù)據(jù)按照x軸的值排序，選出其中的中點(diǎn)進(jìn)行一分為二。在這根線左右兩側(cè)的點(diǎn)被分成了兩棵

5、子樹，對(duì)于這兩個(gè)部分的數(shù)據(jù)來(lái)說(shuō)， 我們更換一個(gè)維度，也就是選擇 y軸進(jìn)行劃分。一樣，我們先排序，然后找到 中間的點(diǎn)，再次一分為二。我們可以得到：我們重復(fù)上述過(guò)程，一直將點(diǎn)分到不能分為止，為了能更好地看清楚，我 們對(duì)所有數(shù)據(jù)標(biāo)上坐標(biāo)（并不精確）。如果我們把空間看成是廣義的區(qū)間，那么它和線段樹的原理是一樣的。最 后得到的也是一棵完美二叉樹，因?yàn)槲覀兠看味歼x擇了數(shù)據(jù)集的中點(diǎn)進(jìn)行劃分， 可以保證從樹根到葉子節(jié)點(diǎn)的長(zhǎng)度不會(huì)超過(guò)叩N。我們代入上面的坐標(biāo)之后, 我們最終得到的KD-Tree大概是下面這個(gè)樣子：KD-Tree 建樹在建樹的過(guò)程當(dāng)中，我們的樹深每往下延伸一層，我們就會(huì)換一個(gè)維度作 為衡量標(biāo)準(zhǔn)。原

6、因也很簡(jiǎn)單，因?yàn)槲覀兿Ｍ@棵樹對(duì)于這K維空間都有很好的表達(dá)能力，方便我們根據(jù)不同的維度快速查詢。在一些實(shí)現(xiàn)當(dāng)中，我們會(huì)計(jì)算 每一個(gè)維度的方差，然后選擇方差較大的維度進(jìn)行切分。這樣做自然是因?yàn)榉?差較大的維度說(shuō)明數(shù)據(jù)相對(duì)分散，切分之后可以把數(shù)據(jù)區(qū)分得更加明顯。但我 個(gè)人覺(jué)得這樣做意義不是很大，畢竟計(jì)算方差也是一筆開(kāi)銷。所以這里我們選 擇了最樸素的方法一一輪流選擇。也就是說(shuō)我們從樹根開(kāi)始，選擇第0維作為排序和切分?jǐn)?shù)據(jù)的依據(jù)，然后到了樹深為1的這一層，我們選擇第一維，樹深為2的這一層，我們選擇第二維，以此類推。當(dāng)樹深超過(guò)了K的時(shí)候，我們就對(duì)樹深取模。明確了這一點(diǎn)之后，我們就可以來(lái)寫KD-Tree的建

7、樹代碼了，和上面二叉樹的代碼非常相似，只不過(guò)多了維度的處理而已。#外部暴露接口def build_model(self, dataset):self.root = self._build_model(dataset)#先忽略，容后再講self.set_father(self.root, None)#內(nèi)部實(shí)現(xiàn)的接口def _build_model(self, dataset, dep th=O): if len( dataset) = 0:return None#通過(guò)樹深對(duì)K取模來(lái)獲得當(dāng)前對(duì)哪一維切分axis = depth % self.Km = len( dataset) / 2#根據(jù)axi

8、s這一維排序dataset = sorted(dataset, key=lambda x: xaxis)#將數(shù)據(jù)一分為二left, mid, right = dataset:m, datasetm, datasetm+1:#遞歸建樹return KDTree.Node(midaxis,mid,axis,dep th,len( dataset),self._build_model(left, dep th+1), self._build_model(nght, dep th+1)這樣我們就建好了樹，但是在后序的查詢當(dāng)中我們需要訪問(wèn)節(jié)點(diǎn)的父節(jié)點(diǎn)， 所以我們需要為每一個(gè)節(jié)點(diǎn)都賦值指向父親節(jié)點(diǎn)的指針。

9、這個(gè)值我們可以寫在 建樹的代碼里，但是會(huì)稍稍復(fù)雜一些，所以我把它單獨(dú)拆分了出來(lái)，作為一個(gè) 獨(dú)立的函數(shù)來(lái)給每一個(gè)節(jié)點(diǎn)賦值。對(duì)于根節(jié)點(diǎn)來(lái)說(shuō)，由于它沒(méi)有父親節(jié)點(diǎn)，所 以賦值為Nones我們來(lái)看下set_father當(dāng)中的內(nèi)容，其實(shí)很簡(jiǎn)單，就是一個(gè) 樹的遞歸遍歷：def set_father(self, no de, father):if node is None:return#賦值no de.father = father#遞歸左右self.set_father( no de.lchild, no de)self.set_father( no de.rchild, no de)快速批量查詢KD-Tr

10、ee建樹建好了肯定是要來(lái)用的，它最大的用處是可以在單次查詢中 獲得距離樣本最近的若干個(gè)樣本。在分散均勻的數(shù)據(jù)集當(dāng)中，我們可以在的時(shí) 間內(nèi)完成查詢，但是對(duì)于特殊情況可能會(huì)長(zhǎng)一些，但是也比我們通過(guò)樸素的方 法查詢要快得多。我們很容易發(fā)現(xiàn)，KD-Tree 一個(gè)廣泛的使用場(chǎng)景是用來(lái)優(yōu)化KNN算法。我們?cè)谥敖榻BKNN算法的文章當(dāng)中曾經(jīng)提到過(guò)，KNN算法在預(yù)測(cè)的時(shí)候需要遍 歷整個(gè)數(shù)據(jù)集，然后計(jì)算數(shù)據(jù)集中每一個(gè)樣本與當(dāng)前樣本的距離，選出最近的 K個(gè)來(lái)，這需要大量的開(kāi)銷。而使用KD-Tree，我們可以在一次查詢當(dāng)中直接查找到K個(gè)最近的樣本，因此大大提升 KNN算法的效率。那么，這個(gè)查詢操作又 是怎么實(shí)現(xiàn)的呢

11、？這個(gè)查詢基于遞歸實(shí)現(xiàn)，因此對(duì)于遞歸不熟悉的小伙伴，可 能初看會(huì)比較困難，可以先閱讀一下之前關(guān)于遞歸的文章。首先我們先通過(guò)遞歸查找到KD-Tree上的葉子節(jié)點(diǎn)，也就是找到樣本所在 的子空間。這個(gè)查找應(yīng)該非常容易，本質(zhì)上來(lái)說(shuō)我們就是將當(dāng)前樣本不停地與 分割線進(jìn)行比較，看看是在分割線的左側(cè)還是右側(cè)。和二叉搜索樹的元素查找 是一樣的：def iter_dow n( self, no de, data):#如果是葉子節(jié)點(diǎn)，則返回if no de.lchild is None and no de.rchild is None:return node#如果左節(jié)點(diǎn)為空，則遞歸右節(jié)點(diǎn)if no de.lchi

12、ld is None:return self.iter_dow n(no de.rchild, data)#同理，遞歸左節(jié)點(diǎn)if no de.rchild is None:return self.iter_dow n(no de.rchild, data)#都不為空則和分割線判斷是左還是右axis = no de.axisn ext_ node = no de.lchild if dataaxis = no de.bo un dray else no de.rchild return self.iter_dow n(n ext_ no de, data)我們找到了葉子節(jié)點(diǎn)，其實(shí)代表樣本空間當(dāng)中

13、的一小塊空間。我們來(lái)實(shí)際 走一下整個(gè)流程，假設(shè)我們要查找3個(gè)點(diǎn)。首先，我們會(huì)創(chuàng)建一個(gè)候選集，用來(lái)存儲(chǔ)答案。當(dāng)我們找到葉子節(jié)點(diǎn)之后，這個(gè)區(qū)域當(dāng)中只有一個(gè)點(diǎn)，我們把它 加入候選集。在上圖當(dāng)中紫色的x代表我們查找的樣本，我們查找到的葉子節(jié)點(diǎn)之后， 在兩種情況下我們會(huì)把當(dāng)前點(diǎn)加入候選集。第一種情況是候選集還有空余，也 就是還沒(méi)有滿K個(gè)，這里的K是我們查詢的數(shù)量，也就是3。第二種情況是當(dāng) 前點(diǎn)到樣本的距離小于候選集中最大的一個(gè)，那么我們需要更新候選集。這個(gè) 點(diǎn)被我們?cè)L問(wèn)過(guò)之后，我們會(huì)打上標(biāo)記，表示這個(gè)點(diǎn)已經(jīng)訪問(wèn)過(guò)了。這個(gè)時(shí)候 我們需要判斷，整棵樹當(dāng)中的搜索是否已經(jīng)結(jié)束，如果當(dāng)前節(jié)點(diǎn)已經(jīng)是根節(jié)點(diǎn) 了，說(shuō)明

14、我們的遍歷結(jié)束了，那么返回候選集，否則說(shuō)明還沒(méi)有，我們需要繼 續(xù)搜索。上圖當(dāng)中我們用綠色表示樣本被放入了候選集當(dāng)中，黃色表示已經(jīng)訪 問(wèn)過(guò)。由于我們的搜索還沒(méi)有結(jié)束，所以需要繼續(xù)搜索。繼續(xù)搜索需要判斷樣 本和當(dāng)前分割線的距離來(lái)判斷和分割線的另一側(cè)有沒(méi)有可能存在答案。由于葉 子節(jié)點(diǎn)沒(méi)有另一側(cè)，所以作罷，我們往上移動(dòng)一個(gè)，跳轉(zhuǎn)到它的父親節(jié)點(diǎn)。此時(shí)候選集未滿，我們加入候選集，標(biāo)記 但是也沒(méi)有另一側(cè)的節(jié)點(diǎn)，所以也跳過(guò)。 我們執(zhí)行同樣的判斷，發(fā)現(xiàn)此時(shí)候選集還我們計(jì)算距離并且查看候選集, 為已經(jīng)訪問(wèn)過(guò)。它雖然存在分割線, 我們?cè)偻希闅v到它的父親節(jié)點(diǎn), 有空余，于是將它繼續(xù)加入答案：但是當(dāng)我們判斷到分割線

15、距離的時(shí)候，我們發(fā)現(xiàn)這一次樣本到分割線的舉 例要比之前候選集當(dāng)中的最大距離要小，所以分割線的另一側(cè)很有可能存在答 案：這里的d1是樣本到分割線的距離，d2是樣本到候選集當(dāng)中最遠(yuǎn)點(diǎn)的距離。 由于到分割線更近，所以分割線的另一側(cè)很有可能也存在答案，這個(gè)時(shí)候我們 需要搜索分割線另一側(cè)的子樹，一直搜索到葉子節(jié)點(diǎn)。我們找到了葉子節(jié)點(diǎn)，計(jì)算距離，發(fā)現(xiàn)此時(shí)候選集已經(jīng)滿了，并且它的距 離大于候選集當(dāng)中任何一個(gè)答案，所以不能構(gòu)成新的答案。于是我們只是標(biāo)記 它已經(jīng)訪問(wèn)過(guò)，并不會(huì)加入候選集。同樣，我們繼續(xù)往上遍歷，到它的父節(jié)點(diǎn):J比較之后發(fā)現(xiàn)，data到它的距離小于候選集當(dāng)中最大的那個(gè)，于是我們更 新候選集，去掉距

16、離大于它的答案。然后我們重復(fù)上述的過(guò)程，直到根節(jié)點(diǎn)為 止。由于后面沒(méi)有更近的點(diǎn)，所以候選集一直沒(méi)有更新，最后上圖當(dāng)中的三個(gè) 打了綠標(biāo)的點(diǎn)就是答案。我們把上面的流程整理一下，就得到了遞歸函數(shù)當(dāng)中 的邏輯，我們用Python寫出來(lái)其實(shí)已經(jīng)和代碼差不多了：def query (no de, data, an swers, K):#判斷node是否已經(jīng)訪問(wèn)過(guò)if no de.visited:#標(biāo)記訪問(wèn)no de.visited = True#計(jì)算data到node中點(diǎn)的距離dis = cal_dis(data, node.point)#如果小于答案中最大值則更新答案if dis max(a nswer

17、s):an swers. up date( node.point)#計(jì)算data到分割線的距離dis = cal_dis(data, no de.s plit)#如果小于最長(zhǎng)距離，說(shuō)明另一側(cè)還可能有答案if dis max(a nswers):#獲取當(dāng)前節(jié)點(diǎn)的兄弟節(jié)點(diǎn)brother = self.get_brother( no de)if brother is not None:#往下搜索到葉子節(jié)點(diǎn)，從葉子節(jié)點(diǎn)開(kāi)始尋找leaf = self.iter_dow n( brother, data)if leaf is not None:retur n self.query(leaf, data,

18、an swers, K)#如果已經(jīng)到了根節(jié)點(diǎn)了，退出if node is root:return an swers#遞歸父親節(jié)點(diǎn)return seif.query (no de.father, data, an swers, K) else:if node is root:return an swersreturn seif.query (no de.father, data, an swers, K)最終寫成的代碼和上面這段并沒(méi)有太多的差別，在得到距離之后和答案當(dāng) 中的最大距離進(jìn)行比較的地方，我們使用了優(yōu)先隊(duì)列。其他地方幾乎都是一樣 的，我也貼上來(lái)給大家感受一下：def _query_ ne

19、arest_k(self, no de, p ath, data, topK, K):#我們用set記錄訪問(wèn)過(guò)的路徑，而不是直接在節(jié)點(diǎn)上標(biāo)記if node not in p ath:p ath.add( no de)#計(jì)算歐氏距離dis = KDTree.dista nce( no de.value, data)if (len (to pK) K or dis top K-1dista nee):top K.a ppen d( no de: no de, dista nee: dis)#使用優(yōu)先隊(duì)列獲取topKtopK = hea pq.n smallest(K, topK, key=lamb

20、da x: xdista nee)axis = no de.axis#分割線都是直線，直接計(jì)算坐標(biāo)差dis = abs(dataaxis - no de.bo un dray)if len( to pK) K or dis = top K-1dista nee:brother = self.get_brother( no de, p ath)if brother is not None:n ext_ node = self.iter_dow n( brother, data)if n ext_ node is not None:retur n self._query_ nearest_k (n

21、 ext_ no de, p ath, data, topK, K)if node = self.root:return topKretu rn self._query_ nearest_k (no de.father, p ath, data, topK, K) else:if node = self.root:return topKretu rn self._query_ nearest_k (no de.father, p ath, data, topK, K)這段邏輯大家應(yīng)該都能看明白，但是有一個(gè)疑問(wèn)是，我們?yōu)槭裁床辉趎ode里面加一個(gè)visited的字段，而是通過(guò)傳入一個(gè)set來(lái)維護(hù)訪問(wèn)過(guò)的節(jié)點(diǎn)呢？這個(gè)邏輯只看代碼是很難想清楚的，必須要親手實(shí)驗(yàn)才會(huì)理解。如果在node當(dāng)中加入一個(gè)字段當(dāng)然也是可以的，如