[Data Structure] K-D Tree

What is K-D Tree?

K-D tree는 공간을 기반으로 partitioning하는 binary tree인데, 점들을 k-차원 공간에서 정리한다. 

wikipedia k-d tree

 위의 그림은 3차원 k-d 트리의 예시인데, 

1) 빨간색의 수직 plane으로 split하여, 하얀색 선으로 표시된 cell을 두개의 subcell로 나눈다. 

2) 각자의 subcell은 초록색의 수평 plane으로 split되어, 또 두 개의 subcell들로 나뉜다.

3) 네 개의 cell은 파란색의 수직 plane들로 각각 split되어, 두 개의 subcell들로 나뉜다.

이렇게, splitting이 끝난 후의 마지막 8개의 cell들을 leaf cell이라고 한다. 

 

이러한 k-d tree는 보통 다차원 공간에서의 search 등에 많이 사용되는 기법이며, range search, nearest neighbor 등의 문제에 적용된다.

• range search : 2차원, 3차원 물체가 서로 겹치는지 확인하여, 겹치는 물체들을 모두 search하는 것.

삼각형(range) 안의 search

• neareast neighbor : 주어진 point와 가장 가까운 (주어진 set안의) 점을 찾는 것.

K-D Tree는 

 

Implementation

from collections import namedtuple
from operator import itemgetter
from pprint import pformat

class Node(namedtuple("Node", "location left_child right_child")):
    def __repr__(self):
        return pformat(tuple(self))

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

    k = len(point_list[0])  # assumes all points have the same dimension
    # Select axis based on depth so that axis cycles through all valid values
    axis = depth % k

    # Sort point list by axis and choose median as pivot element
    point_list.sort(key=itemgetter(axis))
    median = len(point_list) // 2

    # Create node and construct subtrees
    return Node(
        location=point_list[median],
        left_child=kdtree(point_list[:median], depth + 1),
        right_child=kdtree(point_list[median + 1 :], depth + 1),
    )

def main():
    """Example usage"""
    point_list = [(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)]
    tree = kdtree(point_list)
    print(tree)

if __name__ == "__main__":
    main()

위의 코드는 https://en.wikipedia.org/wiki/K-d_tree 에 기재된 python 구현 코드이다.

(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)의 6개의 코드를 K-D Tree를 사용하면, 아래와 같이 출력된다. 

간단하게 설명하자면, 위의 예제에서 각 좌표는 2차원이므로, x축(depth = 0)과 y축(depth = 1)을 기준으로 번갈아가면서 모든 point를 노드에 넣을 때까지 partitioning하는 것이다. 각 축의 순서대로 정렬된 후 median 값을 기준으로 왼쪽과 오른쪽 subtree로 나누는 과정을 재귀적으로 진행한다. 

 

Nearest neighbour search with K-D Tree

Nearest Neighbour search 알고리즘은, tree안의 point들 중 주어진 point와 가장 가까운 것을 찾아내는 것이다. 

K-D Tree안의 nearest neighbour를 찾는 순서는 아래와 같다.

1. 루트 노드부터 시작해서 tree를 recursively하게 아래로 내려가면서, search point가 이 tree에 insert된다면 들어가야 할 자리를 찾음. 
2. 만약 리프 노드에 도착하면, 노드 point를 체크한 후, distance가 더 가까우면 이 노드 point를 "current best"로 저장.
3. tree의 recursion을 거슬러올라가면서, 각 노드마다 아래의 과정을 거친다.  

        (1) 만약 현재 node가 current best 노드보다 더 가까이 있다면, current best를 업데이트함.

        (2) 반대 splitting plane을 넘어서, search point와 더 가까운 점이 있을지 확인함. 

           - search point 중심으로, current nearest distance거리를 반지름으로 한 hypersphere를 그린 후, 이와 교차하는              점이 있는지 확인하면 된다.

             (a) 만약, hypersphere가 plane을 교차할 때, 반대 side에 nearer point가 있을 수도 있으므로, 알고리즘은 현재

              노드에서 다른 branch로 이동하여 똑같은 recursive process를 거쳐야 한다. 

             (b) 만약, hypersphere가 splitting plane과 교차하지 않는다면, 알고리즘은 반대 side의 경우를 전부 제거하고                마저 진행해도 됨. 

   4. 알고리즘이 루트 노드에 대해 위와 같은 process를 끝마치면, search가 끝난다. 

 

https://www.youtube.com/watch?v=Y4ZgLlDfKDg

위의 경우, 새로운 점 (7,4)와 가장 가까운 점을 찾는 예시이다.  (7,4)는 KD-tree를 따라가다보면, 오른쪽 그림에서 파란색으로 칠해져있는 bounding box안에 속하게 된다. (leaf node) 가장 아래의 노드인 (7,2)와 (9,6)과의 distance를 구하고, 제일 작은 distance를 "current best" distance로 설정한다. 이후, 그 distance를 radius로 한 hypersphere를 그렸을 때, 반대 편에 속하는 점이 이 sphere안에 있기 때문에, 다른 branch도 체크하여야 한다.