Поиск по диапазону: от отрезка на прямой к прямоугольнику в kd-дереве

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

Как думаешь, сколько это стоит

Каталог из пяти тысяч товаров, запрос «цена от 250 до 260 тысяч». Отсортируй по цене - и ответ два двоичных поиска. А если запрос двумерный: «цена от 250 до 260 тысяч и рейтинг от 4 до 5»? Отсортировать можно и по двум ключам, но прямоугольник запроса в таком порядке подряд не лежит.

Ответь себе до чтения: сколько узлов kd-дерева из 1024 точек обойдёт запрос-полоса, внутри которой одна точка? Столько же, сколько нашёл, - один-два? Ответ (опыт ниже: 1024 случайные точки, полоса шириной 10): 62 узла. Цена запроса зависит не только от ответа, но и от того, сколько ячеек задевает прямоугольник.

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

Поиск в архиве, разложенном по шкафам: шкафы по годам, полки по месяцам, папки по дням. Запрос «всё за март 2019 года» - открываешь один шкаф и одну полку. Запрос «все документы, подписанные 15-го числа любого месяца» - приходится заглянуть на каждую полку каждого шкафа, хотя документов мало. Разложение по одному ключу плохо помогает запросу по другому.

Механизм: двоичный поиск и ячейки kd-дерева

Одномерный запрос. Отсортированный массив, запрос [lo, hi]. Первый индекс со значением ≥ lo и первый со значением > hi находятся двоичным поиском; между ними - ответ. Подсчёт - разность индексов, O(log n); выдача - ещё k. То же даёт сбалансированное дерево поиска: спуск к lo и к hi, а между путями - целые поддеревья, которые лежат в диапазоне полностью.

Двумерный запрос в kd-дереве. Каждый узел отвечает за ячейку - прямоугольник, где лежат точки его поддерева. Запрос заходит в левое поддерево, если прямоугольник заходит левее разреза, в правое - если правее. Если ячейка целиком внутри прямоугольника, всё поддерево - ответ, и его можно выдать без проверок. Учебная программа ниже так не делает и проверяет каждый узел: для такой ячейки это те же k узлов ответа, поэтому оценка не меняется.

Откуда √n. Посмотрим, сколько ячеек пересекает вертикальная прямая. Корень режет по x - прямая идёт в одну половину. Следующий уровень режет по y - прямая пересекает обе дочерние ячейки. Значит, на каждые два уровня число задетых ячеек удваивается, а точек в каждой становится вчетверо меньше: Q(n) = 2 + 2 · Q(n/4), откуда Q(n) = O(√n). Узлы, чьи ячейки прямоугольник задевает, но не содержит, - это ячейки, пересечённые одной из четырёх сторон прямоугольника, то есть O(√n); остальное - ответ, k.

Правило

На прямой диапазон - два двоичных поиска, O(log n + k). В kd-дереве прямоугольный запрос заходит только в ячейки, которые прямоугольник задевает, и выдаёт целиком ячейки, которые он содержит: O(√n + k) на плоскости, O(n^(1 − 1/d) + k) в d измерениях. Корень - от того, что сторона прямоугольника пересекает O(√n) ячеек.

Разбор: цены на прямой

Пять тысяч цен, три запроса: узкий, весь диапазон и край. Считаем шаги двоичного поиска.

import random
schet = [0]

def pervyj_ne_menshe(a, x):
    # первый индекс i, где a[i] >= x (двоичный поиск)
    lo = 0
    hi = len(a)
    while lo < hi:
        schet[0] = schet[0] + 1
        mid = (lo + hi) // 2
        if a[mid] < x:
            lo = mid + 1
        else:
            hi = mid
    return lo

random.seed(5)
ceny = []
for i in range(5000):
    ceny.append(random.randint(1, 1000000))
ceny.sort()
for lo, hi in [[250000, 260000], [1, 1000000], [999000, 999999]]:
    schet[0] = 0
    k = pervyj_ne_menshe(ceny, hi + 1) - pervyj_ne_menshe(ceny, lo)
    print(lo, hi, k, schet[0])

250000 260000 47 25
1 1000000 5000 25
999000 999999 3 24

Сорок семь товаров, все пять тысяч и три - а шагов поиска одинаково, 24-25: два двоичных поиска по log₂ 5000 ≈ 12. Подсчёт не зависит от ответа. hi + 1 работает, потому что цены целые; для дробных значений правую границу ищут вторым поиском - первый индекс, где значение > hi. Если ответ нужно выдать списком, добавится k - но это цена самого ответа, её не обойти.

Разбор: квадрат, полосы и весь квадрат в kd-дереве

1024 случайные точки в квадрате 10 000 × 10 000, четыре прямоугольника. Печатаем число найденных точек и посещённых узлов.

import random

def postroit(tochki, glubina):
    # узел: [точка, ось, левое, правое]; медиана по оси по очереди
    if len(tochki) == 0:
        return None
    os = glubina % 2
    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 v_pryamougolnike(uzel, pr, otvet, schet):
    # pr = [x1, x2, y1, y2]; собрать точки внутри прямоугольника
    if uzel == None:
        return
    schet[0] = schet[0] + 1
    p, os, levo, pravo = uzel
    if pr[0] <= p[0] <= pr[1] and pr[2] <= p[1] <= pr[3]:
        otvet.append(p)
    lo = pr[2 * os]
    hi = pr[2 * os + 1]
    # в левое поддерево - если прямоугольник заходит левее разреза
    if lo <= p[os]:
        v_pryamougolnike(levo, pr, otvet, schet)
    if hi >= p[os]:
        v_pryamougolnike(pravo, pr, otvet, schet)
random.seed(6)
tochki = []
for i in range(1024):
    tochki.append([random.randint(0, 10000), random.randint(0, 10000)])
derevo = postroit(tochki, 0)
zaprosy = [[4000, 5000, 4000, 5000], [5000, 5010, 0, 10000],
           [0, 10000, 5000, 5010], [0, 10000, 0, 10000]]
for pr in zaprosy:
    otvet = []
    schet = [0]
    v_pryamougolnike(derevo, pr, otvet, schet)
    print(pr, len(otvet), schet[0])

[4000, 5000, 4000, 5000] 10 38
[5000, 5010, 0, 10000] 1 62
[0, 10000, 5000, 5010] 1 93
[0, 10000, 0, 10000] 1024 1024

При равенстве границы и разреза программа заходит в оба поддерева: точки, равные медиане по оси, бывают в обоих. Квадрат 1000 × 1000 - десять точек за 38 узлов. Вертикальная и горизонтальная полосы шириной 10 - по одной точке, но 62 и 93 узла: полоса пересекает ячейки по всей длине, как на рисунке. √1024 = 32, и полосы стоят два-три корня - ровно то, что предсказывает рекуррентность. Весь квадрат - 1024 узла: это ответ, и меньше нельзя.

Граница: √n при миллионе точек - тысяча узлов на запрос, даже пустой. Если запросов миллионы и они узкие, нужна структура с логарифмической ценой, - дерево диапазоновРазбирается в разделе 23, глава «Поиск по диапазону и BSP-деревья»: Дерево диапазонов: логарифм в квадрате на запрос и каскадирование из следующего урока.

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

Проверять каждую точку ячейки, которая целиком внутри, когда нужен только подсчёт. Для выдачи это те же k шагов, а подсчёт большого прямоугольника вырастает до n. Хранить в узле размер поддерева - и ячейка внутри считается одной операцией.

Два одномерных индекса вместо двумерного. База найдёт по индексу цены тысячу товаров и отфильтрует их по рейтингу, хотя ответ - три товара. Пересечение двух списков стоит сумму их длин, а не размер ответа.

kd-дерево в высокой размерности для диапазонов. Оценка n^(1 − 1/d): в десяти измерениях это n^0,9 - почти перебор.

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

Фильтры каталога по двум-трём числовым признакам, запрос «все объекты в окне карты», выборка событий по времени и месту. Базы данных решают это составными и пространственными индексами (GiST, R-деревья), а аналитические системы - деревьями и предрасчётом. Понимание √n объясняет, почему узкие длинные окна работают медленнее, чем квадратные той же площади.

Теперь сам

1. Массив из миллиона отсортированных чисел. Сколько примерно сравнений нужно, чтобы посчитать числа в диапазоне?
2. kd-дерево из n = 10⁶ точек. Порядок числа узлов для пустой узкой полосы?
3. Почему вертикальная прямая пересекает обе дочерние ячейки после разреза по y, но только одну после разреза по x?