kd tree

kd 트리는 이진 트리로, 모든 트리의 노드가 k차원으로 이루어져 있다.

 

최근 30차원의 데이터를 다루며 검색 및 존재 하지 않을 경우 유사도(유클리디안 거리)를 구하기 위해

 

보게 되었는데, 아쉽게도 트리를 구성하는 데이터가 10만개가 넘다보니 오래 걸렸다. (내 컴퓨터가 구린 것일 수도...)

 

어쨌든 간단하게 살펴보면,

 

a(7,2), b(5,4), c(9,6), d(4,7), e(8,1) 이 순으로 데이터가 트리가 삽입된다 하자

 

먼저 k=2가 된다.

 

트리의 구성

a는 트리가 비어 있으므로 루트노드가 된다.

 

a의 x좌표를 비교하여 작은 애들은 왼쪽으로 큰 애들은 오른쪽으로 분류한다.

 

따라서 b는 a의 왼쪽 노드, c는 오른쪽 노드가 된다.

 

d의 x좌표는 a의 x좌표보다 작으므로 왼쪽으로 가지만 이미 b가 있으므로 이번에는 y좌표를 비교한다.

 

b의 y좌표보다 d의 y좌표가 크므로 오른쪽에 위치한다. 동일하게 e도 배치하면 위 그림과 같이 트리가 구성된다.

 

 

데이터 순서에 따라 트리의 구성이 달라지고 그에 따라 검색 속도도 천차만별로 바뀌게 된다.

 

이것을 막고자 트리 구성시 정렬 후 중간값을 취하도록 하는 방법도 있다.

from operator import itemgetter

class Node:
    def __init__(self, data, depth, left_child=None, right_child=None):
        self.left_child = left_child
        self.right_child = right_child
        self.data = data
        self.depth = depth
    def printALL(self):
        print(self.depth, self.data)
        if self.left_child != None:
            self.left_child.printALL()
        if self.right_child != None:
            self.right_child.printALL()


def kdtree(point_list, depth=0):
    if not point_list:
        return None

    k = len(point_list[0])
    axis = depth % k

    point_list.sort(key=itemgetter(axis))
    median = len(point_list) // 2

    return Node(
        data=point_list[median],
        depth=depth,
        left_child=kdtree(point_list[:median], depth + 1),
        right_child=kdtree(point_list[median + 1:], depth + 1)
    )

if __name__ == '__main__':
    point_list = [(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)]
    tree = kdtree(point_list)
    tree.printALL()

위 코드는 위키백과의 것을 알기 쉽게 수정한 것이다.

 

트리 구성시 정렬하여 중간값을 취하고 재귀적으로 왼쪽과 오른쪽을 구하고 있다.

 

또한 python의 scipy라이브러리내 kdtree의 생성과 검색을 제공하고 있다.

from scipy import spatial

if __name__ == '__main__':
    point_list = [(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)]
    tree = spatial.KDTree(point_list)

    print(tree.query((2, 3)))
    print(tree.query((2, 3.5)))

트리의 인자값으로 넣을 컨테이너로써 튜플이나 리스트외에도

 

DataFrame과 ndarray도 넣을 수 있다. query메소드로 검색하면 결과로

 

유클리디안 거리와 가장 유사한 노드의 인덱스 번호를 반환한다.

 

실행 결과