kd-дерево: ближайший сосед без перебора всех точек

Если ты знаешь, как строится kd-дерево, как устроен поиск ближайшего с отсечением по разделяющей прямой и почему в больших размерностях оно вырождается в перебор, листай до «Теперь сам».

Как думаешь, сколько точек посмотреть

Шестьсот банкоматов на карте, и для каждого из сорока запросов «где ближайший?» надо ответить. Перебор - шестьсот расстояний на запрос. Отсортировать банкоматы по x и смотреть только соседей по x нельзя: ближайший может стоять прямо над тобой, а по x - далеко в списке.

Ответь себе до чтения: сколько расстояний придётся посчитать, если разрезать карту пополам, каждую половину ещё пополам и так далее, а при поиске заглядывать в соседний кусок, только если он может содержать что-то ближе найденного? Ответ: около четырнадцати на запрос. А если банкоматы стоят не на плоскости, а точки - восьмимерные векторы признаков, - двести из трёхсот.

На что это похоже

Поиск ближайшей аптеки в незнакомом городе. Сначала определяешь, в какой половине города ты - северной или южной, потом в какой половине этой половины, пока не дойдёшь до своего квартала. Находишь аптеку в квартале. Потом смотришь: граница соседнего квартала ближе, чем найденная аптека? Если да - туда стоит заглянуть, вдруг там аптека прямо за углом. Если граница дальше - в соседний квартал даже не заходишь.

Механизм: разрезы по медиане и отсечение

Построение. Узел хранит точку-медиану и ось разреза. Левое поддерево - точки с меньшей координатой по этой оси, правое - с большей; равные медиане попадают в любую сторону, и отсечение от этого не ломается: до них всё равно не ближе, чем до разделяющей прямой. На следующем уровне ось меняется: x, y, x, y, … (в d измерениях - по кругу). Медиана делит точки пополам, поэтому глубина дерева log₂ n. Построение сортировкой на каждом уровне стоит n · log² n; с выбором медианы за линейное время - n · log n.

Поиск ближайшего.

искать(узел, q):
    проверить точку узла: не ближе ли она лучшего
    сначала - поддерево по ту же сторону разреза, что q
    потом - другое поддерево, но только если
        (расстояние от q до разделяющей прямой)² < лучшее

Почему отсечение верно. Все точки другого поддерева лежат по ту сторону разделяющей прямой. Расстояние от q до любой из них не меньше расстояния от q до самой прямой. Если прямая дальше лучшего найденного, ближе там быть никого не может.

Правило

kd-дерево делит точки медианой по координатам по очереди; глубина log n, построение O(n · log n). Поиск ближайшего спускается к ячейке запроса и возвращается в соседнее поддерево, только если разделяющая плоскость ближе найденного кандидата. В малой размерности на случайных данных это порядка log n узлов, но гарантии логарифма в худшем случае нет, а с ростом d отсечение вырождается в перебор.

Разбор: шестьсот банкоматов

import random

def kv(p, q):
    s = 0
    for i in range(len(p)):
        s = s + (p[i] - q[i]) ** 2
    return s

def postroit(tochki, glubina):
    # узел: [точка, ось, левое поддерево, правое поддерево]
    if len(tochki) == 0:
        return None
    os = glubina % len(tochki[0])
    pary = []
    for p in tochki:
        pary.append([p[os], p])
    pary.sort()
    m = len(pary) // 2
    levye = []
    for para in pary[:m]:
        levye.append(para[1])
    pravye = []
    for para in pary[m + 1:]:
        pravye.append(para[1])
    return [pary[m][1], os, postroit(levye, glubina + 1),
            postroit(pravye, glubina + 1)]

def iskat(uzel, q, luchshee, schet):
    # luchshee = [квадрат расстояния, точка] - лучшее найденное
    if uzel == None:
        return
    schet[0] = schet[0] + 1
    p, os, levo, pravo = uzel
    d = kv(p, q)
    if d < luchshee[0]:
        luchshee[0] = d
        luchshee[1] = p
    raznica = q[os] - p[os]
    blizhnee = levo
    dalnee = pravo
    if raznica > 0:
        blizhnee = pravo
        dalnee = levo
    iskat(blizhnee, q, luchshee, schet)
    # дальнее поддерево - только если разделяющая прямая ближе лучшего
    if raznica * raznica < luchshee[0]:
        iskat(dalnee, q, luchshee, schet)
random.seed(1)
n = 600
tochki = []
for i in range(n):
    tochki.append([random.randint(0, 10000), random.randint(0, 10000)])
derevo = postroit(tochki, 0)
schet = [0]
verno = 0
for t in range(40):
    q = [random.randint(0, 10000), random.randint(0, 10000)]
    luchshee = [10 ** 12, None]
    iskat(derevo, q, luchshee, schet)
    perebor = 10 ** 12
    for p in tochki:
        d = (p[0] - q[0]) ** 2 + (p[1] - q[1]) ** 2
        if d < perebor:
            perebor = d
    if luchshee[0] == perebor:
        verno = verno + 1
print(n, verno, schet[0] / 40)

600 40 13.975

Все 40 ответов совпали с перебором, на запрос в среднем 14 посещённых узлов из 600. Логарифм log₂ 600 ≈ 9,2 - это спуск до листа, остальное - заглядывания в соседние ячейки, которые проверка не смогла отсечь. Расстояния везде квадратами: сравнение с квадратом расстояния до прямой raznica² - точное в целых.

Сравни с перебором: 600 расстояний на запрос. Выигрыш сорокакратный - и на случайных точках плоскости растёт с n: глубина растёт как логарифм, а перебор - линейно.

Разбор: проклятие размерности

Та же структура, триста случайных точек, но в 2, 4 и 8 измерениях. Ось разреза идёт по кругу по всем координатам.

import random

def kv(p, q):
    s = 0
    for i in range(len(p)):
        s = s + (p[i] - q[i]) ** 2
    return s

def postroit(tochki, glubina):
    # узел: [точка, ось, левое поддерево, правое поддерево]
    if len(tochki) == 0:
        return None
    os = glubina % len(tochki[0])
    pary = []
    for p in tochki:
        pary.append([p[os], p])
    pary.sort()
    m = len(pary) // 2
    levye = []
    for para in pary[:m]:
        levye.append(para[1])
    pravye = []
    for para in pary[m + 1:]:
        pravye.append(para[1])
    return [pary[m][1], os, postroit(levye, glubina + 1),
            postroit(pravye, glubina + 1)]

def iskat(uzel, q, luchshee, schet):
    # luchshee = [квадрат расстояния, точка] - лучшее найденное
    if uzel == None:
        return
    schet[0] = schet[0] + 1
    p, os, levo, pravo = uzel
    d = kv(p, q)
    if d < luchshee[0]:
        luchshee[0] = d
        luchshee[1] = p
    raznica = q[os] - p[os]
    blizhnee = levo
    dalnee = pravo
    if raznica > 0:
        blizhnee = pravo
        dalnee = levo
    iskat(blizhnee, q, luchshee, schet)
    # дальнее поддерево - только если разделяющая прямая ближе лучшего
    if raznica * raznica < luchshee[0]:
        iskat(dalnee, q, luchshee, schet)
random.seed(2)
n = 300
for d in [2, 4, 8]:
    tochki = []
    for i in range(n):
        p = []
        for k in range(d):
            p.append(random.randint(0, 1000))
        tochki.append(p)
    derevo = postroit(tochki, 0)
    schet = [0]
    for t in range(15):
        q = []
        for k in range(d):
            q.append(random.randint(0, 1000))
        iskat(derevo, q, [10 ** 12, None], schet)
    print(d, n, round(schet[0] / 15, 1), round(schet[0] / 15 / n, 2))

2 300 13.3 0.04
4 300 37.2 0.12
8 300 200.1 0.67

В двумерии запрос смотрит 4% узлов, в четырёхмерии - 12%, в восьмимерии - две трети. Причина - в геометрии больших размерностей: случайные точки почти равноудалены друг от друга, и расстояние до ближайшего соседа сравнимо с расстоянием до любой разделяющей плоскости. Отсечение «граница дальше кандидата» почти никогда не срабатывает.

Необходимое условие выходит прямо из числа: kd-дерево имеет смысл, только когда точек заметно больше, чем 2^d: чтобы каждую ось разрезать хотя бы раз, нужно d уровней, то есть 2^d ячеек, - а меньше точек столько ячеек не наберут, и разрезы по части осей не случатся вовсе. Условие не достаточное: при d = 8 и 300 точках их больше 2⁸, а просмотрено всё равно 67% дерева - нужно n, во много раз больше 2^d. Начальное «лучшее» 10 ** 12 в программе - просто число больше любого квадрата расстояния: здесь их потолок 2 · 10 000². Для векторов признаков с сотнями измерений его не используют вовсе - там работают приближённые методы, которые ищут не точного ближайшего, а одного из близких: локально-чувствительное хеширование (LSH - хеш-функции, у которых близкие векторы с большой вероятностью попадают в одну корзину) и графы близости вроде HNSW в векторных базах данных.

Что ломается без этого

Отсечение по прямоугольнику, но без ячейки. Распространённая ошибка - проверять расстояние до точки узла, а не до разделяющей прямой: точка узла может быть далеко, а граница - близко, и соседняя половина пропускается вместе с ответом.

Порядок обхода: сначала дальнее поддерево. Ответ не меняется, но лучший кандидат найдётся поздно, и отсечение не сработает почти нигде - время приближается к перебору.

kd-дерево для частых вставок и удалений. Медианы перестают быть медианами, дерево перекашивается. Для изменяющихся данных используют R-деревья или перестраивают kd-дерево целиком время от времени.

Зачем это в работе

Ближайший банкомат, курьер, станция зарядки - на плоскости или в трёх измерениях. Облака точек с лидара: для каждой точки найти соседей, чтобы оценить нормаль поверхности. Цветовая палитра: для каждого пикселя найти ближайший цвет из 256 - это трёхмерный поиск. В Python всё это делает scipy.spatial.KDTree (и его быстрый вариант cKDTree); в базе данных - индекс с оператором <-> в PostGIS, который ищет k ближайших по пространственному индексу.

Теперь сам

1. Точки (2, 7), (5, 4), (9, 6), (4, 2), (8, 1). Какая точка будет корнем kd-дерева и по какой оси первый разрез?
2. Запрос (7, 3), лучший найденный кандидат на расстоянии 2. Корень делит по x = 5. Нужно ли заглядывать в левую половину?
3. Почему в восьмимерии отсечение почти не срабатывает, хотя дерево построено так же?