k-d дерево

Какой город ближе всего к точке, где вы стоите? В справочнике GeoNames 45 227 городов мира от десяти тысяч жителей, и перебрать их все — 45 тысяч расстояний на каждый вопрос. Дерево ниже отвечает, посмотрев на два-три десятка. Коснитесь карты в любом месте — подсвеченные города и есть те, на которые оно посмотрело.

Для площади Республики в Алматы ближайший город, понятно, Алматы, в трёх километрах, и дереву хватило 16 городов из 45 227. Если найти так ближайшего соседа для каждого города мира, в среднем выходит 28,6 города на вопрос. Самый одинокий — Нуук: до ближайшего города, Кеблавика, 1404 км. А вот с Северного полюса дереву пришлось перебрать 1115 городов, прежде чем оно убедилось, что ближе всех Хаммерфест, до которого 2150 км. Почему вопросы бывают дешёвыми и дорогими — ниже.

Идея

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

Это k-d дерево — k-мерное дерево поиска; в нашем случае k = 2. Это прямое обобщение двоичного дерева поиска: там ключ один и каждый узел делит прямую на «меньше» и «больше», здесь ключей два, и узлы по очереди делят плоскость то по одному, то по другому. Медиана в каждом узле делает дерево сбалансированным: половины всегда равны, и глубина — log₂ n.

Как ищет ближайшего

Ниже 237 городов вокруг Алматы — от Самарканда до Кашгара — и их дерево: каждая линия делит свой прямоугольник, тот, что достался узлу от предков. Нажмите «играть» и посмотрите, как идёт поиск, или поставьте точку запроса в другом месте.

Поиск начинается с корня, как в обычном дереве поиска: на каждом узле считаем расстояние до его города, запоминаем лучший и спускаемся в ту половину, где лежит запрос. Дойдя до листа, у нас уже есть кандидат, обычно неплохой. Дальше — главное. Поднимаясь, на каждом узле спрашиваем: может ли за его линией оказаться город ближе лучшего? Всё, что за линией, не ближе к запросу, чем сама линия. Значит, если расстояние от запроса до линии больше расстояния до лучшего города, ту сторону можно отбросить целиком, со всеми её городами. Если меньше — придётся заглянуть.

Для площади Республики поиск посмотрел на 8 городов из 237 — почти ровно глубина дерева. Чем дальше запрос от всех городов, тем больше радиус лучшего и тем чаще его круг задевает линии: для точки 42,0° с. ш., 66,5° в. д. ближайший город, Шардара, в 145 км, и смотреть пришлось на 14 городов. Северный полюс в первой сцене — крайний случай того же: лучший город в двух тысячах километров, и почти каждая линия рядом.

Так же ищут k ближайших: вместо одного лучшего держат кучу из k лучших с худшим из них на вершине и сравнивают расстояние до линии с этим худшим. Переключите k в первой сцене.

Земля круглая: третье измерение

У карты мира есть подвох: долгота и широта — не плоские координаты. Градус долготы у экватора — 111 км, а у полюса — ноль, и город «рядом по долготе» может быть за тысячу километров. Дерево на (долгота, широта) находило бы не ближайший город. Выход — подняться на измерение выше. Точку на сфере запишем единичным вектором (cos φ cos λ, cos φ sin λ, sin φ) и будем искать ближайшего по прямой сквозь Землю. Хорда c растёт вместе с дугой: дуга равна 2R·arcsin(c/2R), так что ближайший по хорде — это и ближайший по поверхности. А хорда — обычное евклидово расстояние в трёх измерениях, и k-d дерево с k = 3 подходит без изменений: оси чередуются x, y, z.

Код

Двумерное дерево на 237 городах вокруг Алматы: построение, ближайший город, три ближайших и проверка перебором. Координаты — целые дециметры на плоской карте, поэтому все расстояния точные.

А это первая сцена в цифрах: трёхмерное дерево над единичными векторами всех 45 227 городов, ближайший сосед для каждого из них и несколько точек на краю света.

Сколько это стоит

Построение: на каждом уровне находим медиану и делим точки. С сортировкой на каждом уровне, как в листингах, это O(n log² n); если искать медиану быстрым выбором за линейное время — O(n log n). Памяти — по узлу на точку, O(n).

Поиск ближайшего для точек, разбросанных без особой структуры, в среднем обходит O(log n) узлов: это показали Джером Фридман, Джон Бентли и Рафаэль Финкель в 1977 году, а сам Бентли в первой статье о дереве оценил это по опытам. Гарантии в худшем случае нет: далёкий запрос или точки, сбившиеся на одну прямую, заставляют обойти большую часть дерева. Помогает хранить у каждого узла прямоугольник, в котором лежит всё его поддерево, и сравнивать с лучшим расстояние до прямоугольника, а не до линии, — такую проверку предлагали уже Фридман, Бентли и Финкель, и так сделано в задаче ниже.

Перебор растёт с n линейно, дерево — едва заметно: на 64 тысячах точек оно быстрее в сотни раз на Python и почти в сто раз на Go. Только не забывайте о построении: оно дороже одного перебора, и если запрос один, а точки каждый раз новые, дерево не окупится.

Проклятие размерности

Отсечение работает, пока круг вокруг запроса задевает немного линий. В многомерном пространстве это перестаёт быть так: почти все точки примерно на одном расстоянии от запроса, «шар лучшего» огромен по сравнению с ячейками, и линий он задевает почти все. Листинг строит дерево на 4096 случайных точках в кубе разной размерности и считает, на сколько точек смотрит поиск:

В двух измерениях — 17 точек из 4096, в восьми — уже 747, а в двенадцати — 88% всех точек: дерево уже не быстрее перебора, а обход сверху ещё и стоит времени. Отсюда известное правило: k-d дерево полезно, пока точек намного больше, чем 2^k, где k — размерность. 4096 — это как раз 2¹². Для векторов в сотни измерений — эмбеддингов текстов и картинок — ищут не точно ближайшего, а «достаточно близкого», другими структурами.

Где работает

• Научные библиотеки. scipy.spatial.KDTree в SciPy, поиск соседей в scikit-learn — для точек в двух, трёх, десятке измерений.
• Трассировка лучей. k-d дерево над треугольниками сцены позволяет проверять пересечение луча не со всеми треугольниками, а с теми немногими, что лежат по пути.
• Карты. Ближайший город, станция, адрес к точке на карте; соседи для подсказок «рядом с вами». Для запросов «что попало в окно карты» чаще берут R-деревья или кривые, заполняющие плоскость.
• Классификация по соседям. Метод k ближайших соседей и кластеризация по плотности спрашивают «кто рядом» для каждой точки данных — с деревом это быстро, пока признаков немного.

История

Джон Бентли описал k-d дерево в 1975 году в Communications of the ACM, в статье «Многомерные двоичные деревья поиска для ассоциативного поиска». Он показал, что одна и та же структура отвечает на разные запросы — например, по части ключей и поиск ближайшего соседа, — а для поиска ближайшего заметил по опытам, что время растёт как логарифм числа точек. Через два года вместе с Джеромом Фридманом и Рафаэлем Финкелем он описал поиск ближайших соседей с логарифмическим ожидаемым временем. В той статье ось деления в каждом узле выбирают ту, вдоль которой у точек наибольший разброс, а за линию заглядывают, только если шар вокруг запроса задевает прямоугольник поддерева.

Задачи

В задаче нужно найти k ближайших городов для тысяч запросов, с ничьими и с городами, выстроенными в линию, — там, где простое отсечение по линии проваливается. Один из тестов — те же 45 тысяч городов мира.