KD-Tree로 저차원 최근접 이웃 검색 설계하기
KD-Tree의 축 정렬 공간 분할, 최근접 이웃 가지치기, 재구축 전략과 저차원 데이터에서의 인덱스 선택 기준을 다룬다.
2026-08-15 · 최초 발행 2024-04-29
점 집합을 분할해 검색 후보를 줄이는 구조
K-D Tree(k-dimensional tree)는 k차원 점 집합을 축 정렬 초평면으로 재귀 분할하는 이진 공간 색인이다. 각 노드는 분할에 쓸 차원과 분할 값, 필요하다면 서브트리 경계(bounding box)를 보관한다. 2~3차원 기하 데이터, 로보틱스, GIS, 게임 엔진, 저차원 벡터 검색에서 특히 활용도가 높다.
트리를 만들 때는 분할 축을 라운드로빈 방식이나 분산이 가장 큰 축으로 고르고, 중앙값 근사로 좌우 서브셋의 균형을 맞춘다. 평균 구축 시간은 O(n log n), 메모리 사용량은 O(n)이다.
검색은 질의점이 속한 분할 공간을 따라 후보를 먼저 찾는 것으로 시작한다. 이후 쿼리 구(sphere) 또는 상자와 서브트리 경계 사이의 거리를 계산해 가능성이 없는 영역을 제외한다. 저차원에서는 평균 O(log n) 내의 방문 노드 수를 기대할 수 있지만, 차원이 높아지면 차원의 저주로 인해 성능이 저하된다.
분할 축과 균형이 탐색 비용을 좌우한다
축에 직교하는 초평면으로 공간을 나누면 구현이 단순하고 캐시 지역성도 확보하기 쉽다. 라운드로빈은 단순하고 빠르며, 분산 최대 축을 고르는 방식은 균형과 pruning 효율을 높이는 데 쓸 수 있다. 데이터 분포가 치우치면 트리 균형도 나빠질 수 있으므로 분산 기반 축 선택이 보완책이 된다.
중앙값을 nth_element로 선택하면 균형 잡힌 분할을 만들 수 있다. 선형시간 선택 알고리즘을 적용해 중앙값 선택을 최적화할 수도 있다. 삽입과 삭제를 계속 처리하기보다 배치로 구축하는 편이 적합하며, 점진적 갱신이 필요한 경우에는 재균형 비용을 감안해 주기적 재구축을 운영 전략으로 둔다.
최근접·반경·범위 질의의 가지치기
K-D Tree는 k-NN, r-반경, 직사각형 범위(range/box) 쿼리를 지원한다. 탐색 중에는 서브트리 경계와 쿼리 구 또는 상자의 거리를 비교해 방문할 필요가 없는 서브트리를 제거한다.
대략 d ≤ 20인 저차원에서는 이 방식이 효과적이다. d가 증가하면 방문해야 할 노드가 급증하고 성능은 O(n)에 가까워진다. 거리 계산의 안정성을 위해 float64를 사용하고 특성 스케일을 정규화해야 한다.
중복점(identical points), 공선 또는 공평면 데이터는 분할을 불안정하게 만들 수 있다. 인덱스를 tie-breaker로 사용하거나 jitter와 leaf bucket을 두는 방식으로 대응할 수 있다.
공간 데이터와 저차원 특징 검색에서의 활용
로보틱스와 SLAM에서는 2D/3D 포인트클라우드의 최근접 대응점 탐색(ICP, NDT)을 가속하고, 실시간 obstacle 회피를 위한 kNN 질의에 사용한다. 센서 융합 전처리에서는 노이즈 제거와 다운샘플링을 위한 반경 검색을 적용할 수 있다.
게임 엔진과 물리 시뮬레이션에서는 브로드페이즈 충돌 후보를 뽑는 범위·반경 검색, LOD 전환, AI 에이전트의 근접 탐색에 적합하다. GIS와 지도 서비스에서는 좌표 기반 POI 인접 검색과 뷰포트 범위 쿼리를 처리하며, 캐시 친화적 엣지 배치는 모바일 단말의 저지연 응답에도 활용된다.
추천과 시계열 이상탐지에서는 5~20차원 통계 특징으로 kNN 기반 유사 아이템 탐색과 밀도 추정(LOF) 전처리를 수행할 수 있다. 이 경우에는 배치 주기에 맞춘 재구축 운영이 필요하다.
데이터 특성에 맞춰 인덱스를 고르기
| 구조 | 성능(저차원 kNN) | 확장성(고차원) | 일관성(정확도) | 안정성(업데이트) | 운영 편의 |
|---|---|---|---|---|---|
| K-D Tree | 매우 우수 | 차원↑ 시 급격 저하 | 정확(정밀 검색) | 재구축 필요 | 광범위 라이브러리 |
| Ball Tree | 우수(비축정렬에 강함) | K-D보다 완만 저하 | 정확 | 중간 | 주요 라이브러리 |
| R-Tree | 범위/박스 쿼리 강점 | 중간 | 정확 | 삽입·삭제 용이 | DB/공간엔진 내장 |
| HNSW/Annoy(근사) | 매우 우수(대규모) | 고차원 강함 | 근사 | 동적 삽입 용이 | 운영 성숙 |
| Brute Force | 매우 느림 | 차원 무관 | 정확 | 단순 | 구현 용이(작은 n) |
d ≤ 20이고 n이 크며 정밀한 최근접 이웃 검색이 필요하면 K-D Tree를 우선 검토할 수 있다. d > 30이거나 비유클리드 거리, 빈번한 업데이트가 핵심이라면 Ball Tree, HNSW, R-Tree를 고려한다.
균형 K-D Tree에서 n=1,000,000, d=3일 때 평균 방문 노드는 O(수십수백)이다. 거리 계산은 1e6에서 1e2 수준으로 감소하며, 하드웨어와 분포에 따라 10^310^4배 연산 절감이 가능하다. 저차원에서는 선형 탐색보다 질의 지연을 5~50배 줄일 수 있으며, 편차는 데이터 분포와 하드웨어에 따라 달라진다.
포인트 외 노드와 경계를 포함한 메모리 오버헤드는 1.22.5배다. 5배 가속할 수 있다.leaf_size를 키우면 메모리는 줄지만 단일 leaf의 선형 탐색 비용은 증가한다. 배치 쿼리에서는 캐시 지역성이 높아지고 SIMD·멀티스레드를 적용하면 추가 2
Python으로 구축과 질의 실행하기
전제조건: Python 3.10+, scikit-learn ≥ 1.3, numpy
# pip install scikit-learn numpy
import numpy as np
from sklearn.neighbors import KDTree
# 데이터 준비: 3차원 점 200,000개
rng = np.random.default_rng(42)
X = rng.normal(size=(200_000, 3)).astype(np.float64)
# 스케일 정규화(권장)
X = (X - X.mean(axis=0)) / X.std(axis=0)
# 트리 구축
tree = KDTree(X, leaf_size=40, metric='euclidean')
# 1) k-최근접 이웃(k=5)
q = rng.normal(size=(1, 3))
dist, idx = tree.query(q, k=5, return_distance=True)
print("kNN index:", idx[0], "dist:", dist[0])
# 2) 반경 검색(r=1.2)
ind_radius = tree.query_radius(q, r=1.2)
print("within radius:", len(ind_radius[0]))
# 3) 축 정렬 박스 범위 검색: [lo, hi] 범위
lo = np.array([-0.5, -0.5, -0.5])
hi = np.array([ 0.5, 0.5, 0.5])
# 단순 필터(트리 박스쿼리 API 부재 시):
mask = np.all((X >= lo) & (X <= hi), axis=1)
box_idx = np.where(mask)[0]
print("box count:", box_idx.shape[0])
경량 구현으로 트리를 구축하고 1-NN을 질의하는 예시다.
# python 3.10, numpy 1.26
from __future__ import annotations
import numpy as np
from typing import Optional, Tuple
class KDNode:
__slots__ = ("idxs", "axis", "split", "left", "right", "bbox_min", "bbox_max")
def __init__(self, idxs, axis=None, split=None, left=None, right=None, bbox_min=None, bbox_max=None):
self.idxs = idxs # 리프일 때 포인트 인덱스
self.axis = axis # 분할 축
self.split = split # 분할 값
self.left = left
self.right = right
self.bbox_min = bbox_min # AABB 경계
self.bbox_max = bbox_max
def build_kdtree(X: np.ndarray, leaf_size: int = 16, depth: int = 0) -> KDNode:
n, d = X.shape
idxs = np.arange(n)
return _build(X, idxs, leaf_size)
def _build(X, idxs, leaf_size) -> KDNode:
pts = X[idxs]
bbox_min = pts.min(axis=0)
bbox_max = pts.max(axis=0)
if len(idxs) <= leaf_size or np.allclose(bbox_min, bbox_max):
node = KDNode(idxs=idxs, bbox_min=bbox_min, bbox_max=bbox_max)
return node
# 분산 최대 축 선택
axis = np.argmax(pts.var(axis=0))
# 중앙값 분할
order = np.argpartition(pts[:, axis], len(idxs)//2)
idxs = idxs[order]
median_idx = len(idxs)//2
split_val = X[idxs[median_idx], axis]
left = _build(X, idxs[:median_idx], leaf_size)
right = _build(X, idxs[median_idx:], leaf_size)
node = KDNode(idxs=None, axis=axis, split=split_val, left=left, right=right,
bbox_min=bbox_min, bbox_max=bbox_max)
return node
def _squared_distance(a: np.ndarray, b: np.ndarray) -> float:
diff = a - b
return float(diff @ diff)
def _dist_point_aabb(q: np.ndarray, mn: np.ndarray, mx: np.ndarray) -> float:
# AABB까지의 제곱거리
clamped = np.clip(q, mn, mx)
return _squared_distance(q, clamped)
def query_knn(X: np.ndarray, root: KDNode, q: np.ndarray, k: int = 1) -> Tuple[np.ndarray, np.ndarray]:
import heapq
best = [] # max-heap by negative distance
def visit(node: KDNode):
if node.left is None and node.right is None:
for i in node.idxs:
d2 = _squared_distance(q, X[i])
if len(best) < k:
heapq.heappush(best, (-d2, i))
elif d2 < -best[0][0]:
heapq.heapreplace(best, (-d2, i))
return
# 우선 서브트리 선택
axis = node.axis
go_left = q[axis] <= node.split
first, second = (node.left, node.right) if go_left else (node.right, node.left)
if first is not None:
visit(first)
# 백트래킹 가지치기
if second is not None:
# 현재 경계까지 최소거리로 가지치기
if len(best) < k:
need = True
else:
bound = -best[0][0]
d2_box = _dist_point_aabb(q, second.bbox_min, second.bbox_max)
need = d2_box < bound
if need:
visit(second)
visit(root)
best.sort() # by negative d2
dists = np.sqrt([-x[0] for x in best])
idxs = np.array([x[1] for x in best])
return idxs, dists
if __name__ == "__main__":
rng = np.random.default_rng(42)
X = rng.normal(size=(10000, 8)).astype(np.float64)
root = build_kdtree(X, leaf_size=32)
q = rng.normal(size=(8,))
idxs, dists = query_knn(X, root, q, k=5)
print("indices:", idxs)
print("dists:", np.round(dists, 4))
scikit-learn의 KDTree를 쓰는 경우는 다음과 같다.
# python 3.10, numpy 1.26, scikit-learn 1.4
import numpy as np
from sklearn.neighbors import KDTree
X = np.random.RandomState(0).randn(100000, 6)
tree = KDTree(X, leaf_size=40, metric='euclidean')
q = np.random.rand(1, 6)
dists, idxs = tree.query(q, k=10, return_distance=True)
print(idxs[0], dists[0])
leaf_size는 20~60 범위에서 탐색한다. kNN 비중이 높으면 작게 두고, 반경과 박스 쿼리를 함께 처리한다면 중간값을 선택한다. metric은 기본 euclidean이며, L1 거리가 필요하면 metric='manhattan'을 사용한다. 대규모 배치 쿼리는 joblib 기반 병렬화 또는 파티션 분산 처리 대상으로 삼을 수 있다.
갱신 환경에서 지켜야 할 절충점
특성 스케일을 정규화하고 float64를 사용하면 수치 안정성을 확보할 수 있다. 중앙값 근사 선택과 라운드로빈·분산축 선택을 함께 사용하면 균형과 성능을 절충할 수 있다. 대량 업데이트 환경에서는 주기적 재구축과 Active/Standby 트리 이중화로 무중단 교체를 구성한다.
leaf_size 확대는 메모리를 절감하는 대신 leaf 내부 선형 탐색을 늘린다. 분산 최대 축 선택은 구축 시간을 증가시키지만 pruning 효율을 높인다. 저차원 정밀 검색에는 K-D Tree를, 고차원 대규모 근사 검색에는 HNSW·Annoy·FAISS를 선택하는 구분이 필요하다.