Эта статья полностью посвящена KD-Деревьям: я описываю тонкости построения KD-Деревьев, тонкости реализации функций поиска 'ближнего' в KD-Дереве, а также возможные 'подводные камни', которые возникают в процессе решения тех или иных подзадач алгоритма. Дабы не запутывать читателя терминологией(плоскость, гипер-плоскость и т.п), да и вообще для удобства, полагается что основное действо разворачивается в трехмерном пространстве. Однако же, где нужно я отмечаю, что мы работаем в пространстве другой размерности. По моему мнению статья будет полезна как программистам, так и всем тем, кто заинтересован в изучении алгоритмов: кто-то найдет для себя что-то новое, а кто-то просто повторит материал и возможно, в комментариях дополнит статью. В любом случае, прошу всех под кат.

Введение

KD-Tree(K-мерное дерево), специальная 'геометрическая' структура данных, которая позволяет разбить K-мерное пространство на 'меньшие части', посредством сечения этого самого пространства гиперплоскостями(K > 3), плоскостями (K = 3), прямыми (K = 2) ну, и в случае одномерного пространства-точкой (выполняя поиск в таком дереве, получаем что-то похожее на binary search).
Логично, что такое разбиение обычно используют для сужения диапазона поиска в K-мерном пространстве. Например, поиск ближнего объекта (вершины, сферы, треугольника и т.д.) к точке, проецирование точек на 3D сетку, трассировка лучей (активно используется в Ray Tracing) и т.п. При этом объекты пространства помещаются в специальные параллелепипеды-bounding box-ы(bounding box-ом назовем такой параллелепипед, который описывает исходное множество объектов или сам объект, если мы строим bounding box лишь для него. У точек в качестве bounding box-а берется bounding box с нулевой площадью поверхности и нулевым объемом), стороны которых параллельны осям координат.

Процесс деления узла

Итак, прежде чем использовать KD-Дерево, нужно его построить. Все объекты помещаются в один большой bounding box, описывающий объекты исходного множества (при этом каждый объект ограничен своим bounding box-ом), который затем разбивается (делится) плоскостью, параллельной одной из его сторон, на два. В дерево добавляются два новых узла. Таким же образом каждый из полученных параллелепипедов разбивается на два и т.д. Процесс завершается либо по специальному критерию (см. SAH), либо при достижении определенной глубины дерева, либо же при достижении определенного числа элементов внутри узла дерева. Некоторые элементы могут одновременно входить как в левый, так и в правый узлы (например, когда в качестве элементов дерева рассматриваются треугольники).

Проиллюстрирую данный процесс на примере множества треугольников в 2D:


На рис.1 изображено исходное множество треугольников. Каждый треугольник помещен в свой bounding box, а все множество треугольников вписано в один большой корневой bounding box.


На рис.2 делим корневой узел на два узла (по OX): в левый узел попадают треугольники 1, 2, 5, а в правый-3, 4, 5.


На рис.3, полученные узлы снова делятся на два, а треугольник 5 входит в каждый из них. Когда процесс заканчивается, получаем 4 листовых узла.

Огромное значение имеет выбор плоскости для разделения узла дерева. Существует огромное количество способов сделать это, я привожу лишь некоторые, наиболее часто встречаемые на практике способы(предполагается, что начальные объекты помещены в один большой bounding box, а разделение происходит параллельно одной из осей координат):

? Способ 1: Выбрать наибольшую сторону bounding box. Тогда секущая плоскость будет проходить через середину выбранной стороны.

? Способ 2: Рассекать по медиане: отсортируем все примитивы по одной из координат, а медианой назовем элемент (или центр элемента), который находится на средней позиции в отсортированном списке. Секущая плоскость будет проходить через медиану так, что количество элементов слева и справа будет примерно равным.

? Способ 3: Использование чередования сторон при разбиении: на глубине 0 бьем через середину стороны параллельной OX, следующий уровень глубины-через середину стороны параллельной OY, затем по OZ и т.д. Если мы 'прошлись по всем осям', то начинаем процесс заново. Критерии выхода описаны выше.

? Способ 4: Наиболее 'умный' вариант-использовать SAH (Surface Area Heuristic) функцию оценки разделения bounding box. (Об этом подробно будет рассказано ниже). SAH также предоставляет универсальный критерий остановки деления узла.

Способы 1 и 3 хороши, когда речь идет именно о скорости построения дерева (так как найти середину стороны и провести через нее плоскость, отсеивая элементы исходного множества 'налево' и 'направо', тривиально). В то же время, они довольно часто дают неплотное представление разбиения пространства, что может негативно повлиять на основные операции в KD-Дереве (особенно, при прослеживании луча в дереве).

Способ 2 позволяет достигнуть хороших результатов, но требует немалого количества времени, которое тратится на сортировку элементов узла. Как и способы 1, 3, он довольно прост в реализации.

Наибольший же интерес представляет способ с использованием 'умной' SAH эвристики (способ 4).

Bounding box узла дерева делится N (параллельными осям) плоскостями на N + 1 часть (части, как правило, равны) по каждой из сторон (на самом деле, количество плоскостей для каждой стороны можно быть любым, но для простоты и эффективности берут константу). Далее, в возможных местах пересечения плоскости с bounding box-ом вычисляют значение специальной функции: SAH. Деление производится плоскостью с наименьшим значением SAH функции (на рисунке ниже, я предположил, что минимум достигается в SAH[7], следовательно, деление будет производиться именно этой плоскостью (хотя здесь 2D пространство-так что прямой)):



Значение SAH функции для текущей плоскости вычисляется следующим образом:



В своей реализации KD-Дерева я делю каждую сторону на 33 равные части, используя 32 плоскости. Таким образом, по результатам тестов мне удалось найти 'золотую' середину-скорость работы основных функций дерева/скорость построения дерева.

В качестве SAH эвристики я использую функцию, представленную на рисунке выше.

Стоит отметить, что для принятия решения необходим лишь минимум данной функции по всем секущим плоскостям. Следовательно, используя простейшие математические свойства неравенств, а также отбросив умножение на 2 при расчете площади поверхности узла (в 3D)(SAR, SAL, SA), данную формулу можно существенно упростить. В полной же мере расчеты производятся лишь один раз на узел: при оценке критерия выхода из функции деления. Столь простая оптимизация дает существенный прирост скорости построения дерева (x2.5).

Для эффективного вычисления значения SAH функции необходимо уметь быстро определить, сколько узловых элементов находятся справа от данной плоскости, а сколько-слева. Результаты будут неудовлетворительны, если в качестве алгоритма использовать грубый, так называемый brute force подход с квадратичной асимптотикой. Однако же при использовании 'корзиночного' (binned) метода ситуация значительно меняется в лучшую сторону. Описание данного метода приведено ниже:

Предположим, что мы разбиваем сторону bounding box-а на N равных частей (количество плоскостей-(N-1)), bounding box храним парой координат(pointMin, pointMax-см. рис. 1), тогда за один проход по всем элементам узла мы можем для каждой плоскости точно определить, сколько элементов лежит справа, а сколько-слева от нее. Создадим два массива по N элементов в каждом, и проинициализируем нулями. Пусть это будут массивы с именами aHigh и aLow. Последовательно пробегаем по всем элементам узла. Для текущего элемента проверяем, в какую из частей попадают pointMin и pointMax его bounding box-а. Соответственно, получаем два индекса на элемент множества. Пусть это будут индексы с именами iHigh(для pointMax) и iLow(для pointMin). После этого выполним следующее: aHigh[iHigh] += 1, aLow[iLow] +=1.

После прохода по всем элементам, получаем заполненные массивы aHigh и aLow. Для массива aHigh посчитаем частичные-постфикс (suffix) суммы, а для aLow посчитаем частичные-префикс (prefix) суммы (см. рисунок). Получается, что количество элементов справа (и только справа!) от плоскости с индексом i будет равно aLow[i + 1], количество элементов, лежащих слева (и только слева!): aHigh[i], количество же элементов, которые входят как в левый, так и в правый узлы: AllElements-aLow[i + 1]-aHigh[i].

Поставленная задача решена, а иллюстрация этого нехитрого процесса приведена ниже:



Хотелось бы отметить, что получение заранее известного количества элементов слева и справа от 'бьющей' плоскости позволяет нам заранее выделить нужное количество памяти под них (ведь далее, после получения минимума SAH, нам нужно еще раз пройтись по всем элементам и каждый разместить в нужном массиве, (а использование банального push_back(если reserve-не вызывался), приводит к постоянному аллоцированию памяти-весьма затратная операция), что положительно влияет на скорость работы алгоритма построения дерева (x3.3).

Теперь хотелось бы описать подробнее назначение констант, используемых в формуле расчета SAH, а также рассказать о критерии остановки деления данного узла.

Перебирая константы cI и cT, можно добиться более плотной структуры дерева (или же наоборот), жертвуя временем работы алгоритма. Во многих статьях, посвященных в первую очередь построению KD-Дерева для Ray Tracing рендера, авторы используют значения cI = 1., cT = [1; 2]: чем больше значение cT, тем быстрее строится дерево, и тем хуже результаты прослеживания луча в таком дереве.

В своей же реализации я использую дерево для поиска 'ближнего', и, уделив должное внимание полученным результатам тестирования и поиска нужных коэффициентов, можно заметить, что высокие значения cT-дают нам совершенно не плотно заполненные элементами узлы. Чтобы избежать подобной ситуации, мною было решено выставить значение cT в 1., а значение cI протестировать на разных-больших-единицы данных. В итоге-удалось получить довольно плотную структуру дерева, расплатившись значительным ростом времени при построении. На результатах же поиска 'ближнего' это действие отразилось положительно-скорость поиска увеличилась.

Критерий остановки деления узла, довольно прост:



Иными словами: если стоимость прослеживания дочерних узлов больше стоимости прослеживания родительского узла, то деление нужно прекратить.

Теперь, когда мы научились делить узел KD-Дерева, расскажу про исходные случаи, когда количество элементов в узле получается довольно большим, а критерии остановки по количеству элементов уводят алгоритм в бесконечность. Собственно, все внимание на картинку (на примере треугольников в 2D):



Я называю такие ситуации 'веерными' (имеют общую точку соприкосновения, совпадающие объекты, также отношу в эту категорию). Видно, что как бы мы не проводили секущую плоскость, центральная точка так или иначе попадает в один из узлов, а вместе с ней и треугольники, для которых она является общей.

Процесс построения дерева

Мы научились делить узел дерева, теперь осталось применить полученные знания к процессу построения всего дерева. Ниже я привожу описание своей реализации построения KD-Дерева.

Дерево строится от корня. В каждом узле дерева я храню указатели на левое и правое поддеревья, если таковых у узла нет, то он считается листовым (листом т.е). Каждый узел хранит bounding box, который описывает объекты данного узла. В листовых(и только в листовых!) узлах я храню индексы тех объектов, которые входят в данный узел. Однако, в процессе построения память под узлы выделяется порциями (т.е. сразу для K узлов: во-первых, так эффективнее работать с менеджером памяти, во-вторых,-подряд идущие элементы идеальны для кэширования). Такой подход запрещает хранить узлы дерева в векторе, ибо добавление новых элементов в вектор может приводить к реаллоцированию памяти под все существующие элементы в другое место.

Соответственно, указатели на поддеревья теряют всякий смысл. Я использую контейнер типа список (std::list), каждый элемент которого-вектор (std::vector), с заранее заданным размером (константа). Построение дерева выполняю многопоточно (использую Open MP), то есть каждое поддерево строю в отдельном потоке(x4 к скорости). Для копирования индексов в листовой узел идеально подходит использование move семантики (C++11)(+10% к скорости).

Поиск 'ближнего' к точке в KD-Дереве

Итак, дерево построено, перейдем к описанию реализации операции поиска в KD-Дереве.

Предположим, что поиск мы осуществляем в множестве треугольников: дана точка, требуется найти самый ближний к ней треугольник.

Решать поставленную задачу с применением bruteforce невыгодно: для множества из N точек и M треугольников ассимптотика составит O(N * M). Кроме того, хотелось бы, чтобы алгоритм вычислял расстояние от точки до треугольника 'по-честному' и делал это быстро.

Воспользуемся KD-Деревом и выполним следующее:

? Шаг 1. Найдем самый ближний листовой узел к данной точке (в каждом узле, как известно, мы храним bounding box, а расстояние можем смело вычислять до центра((pointMax + pointMin)*0.5) bounding box-а узла) и обозначим его nearestNode.

? Шаг 2. Методом перебора среди всех элементов найденного узла (nearestNode). Полученное расстояние обозначим minDist.

? Шаг 3. Построим сферу с центром в исходной точке и радиусом длины minDist. Проверим, лежит ли эта сфера полностью внутри (то есть без каких либо пересечений сторон bounding box-узла) nearestNode. Если лежит, то ближний элемент найден, если нет, то перейдем к пункту 4.

? Шаг 4. Запустим от корня дерева, поиск ближнего элемента в радиусе: спускаясь вниз по дереву, проверяем, пересекает ли правый или левый узлы (кроме того, узел может лежать полностью внутри сферы или сфера внутри узла...) данная сфера. Если какой-либо узел был пересечен, то аналогичную проверку выполняем и для внутренних узлов этого же узла. Если мы пришли в листовой узел, то выполним переборный поиск ближнего в данном узле и сравним полученный результат с длиной радиуса сферы. Если радиус сферы больше найденного расстояния, обновим длину радиуса сферы вычисленным значением минимума. Дальнейшие спуски по дереву происходят с обновленной длиной радиуса сферы (если мы используем рекурсивный алгоритм, то радиус просто передается в функцию по ссылке, а далее, где необходимо, обновляется).

Понять суть описанного выше алгоритма помогает следующий рисунок:



По рисунку: предположим, мы нашли ближний листовой узел (голубой на рисунке) к данной точке (выделена красным цветом), тогда, выполнив поиск ближнего в узле, получим, что таковым является треугольник с ичёндексом 1, однако, как видно-это не так. Окружность с радиусом R пересекает смежный узел, следовательно, поиск нужно выполнить и в этом узле, а затем сравнить вновь найденный минимум с тем, что уже есть. Как итог, становится очевидным, что ближним является треугольник с индексом 2.

Теперь мне бы хотелось рассказать об эффективной реализации промежуточных операций, используемых в алгоритме поиска 'ближнего'.

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

Находим проекцию точки A (точка, к которой мы и ищем ближний) на плоскости треугольника. Найденную точку обозначим P. Если P лежит внутри треугольника, то расстояние от A до треугольника равно длине отрезка AP, иначе найдем расстояния от A до каждой из сторон (отрезков) треугольника, выберем минимум. Задача решена.

Описанный алгоритм является не самым эффективным. Более эффективный подход опирается на поиск и анализ (нахождение минимума градиента и т.п.) функции, значениями которой являются все возможные расстояния от данной точки до любой точки в треугольнике. Метод довольно сложен в восприятии и, как я считаю, заслуживает отдельной статьи (пока что он реализован у меня в коде, а ссылку на код вы найдете ниже). Ознакомиться с методом можно по [1]. Данный метод по результатам тестов оказался в 10 раз быстрее того, что я описывал ранее.

Определить находится ли сфера с центром в точке O и радиусом R внутри конкретного узла, представленного bounding box, просто(3D):

inline bool isSphereInBBox(const SBBox& bBox, const D3& point, const double& radius)
{
    return (bBox.m_minBB[0] < point[0] - radius && bBox.m_maxBB[0] > point[0] + radius &&
                bBox.m_minBB[1] < point[1] - radius && bBox.m_maxBB[1] > point[1] + radius &&
                bBox.m_minBB[2] < point[2] - radius && bBox.m_maxBB[2] > point[2] + radius);
}

С алгоритмом же определения пересечения сферы с bounding box-ом узла, нахождения узла внутри сферы или сферы внутри узла, дела обстоят несколько иначе. Снова проиллюстрирую (картинка взята с [2]) и приведу корректный код, позволяющий выполнить данную процедуру (в 2D, 3D-аналогично):

bool intersects(CircleType circle, RectType rect)
{
    circleDistance.x = abs(circle.x - rect.x);
    circleDistance.y = abs(circle.y - rect.y);

    if (circleDistance.x > (rect.width/2 + circle.r)) { return false; }
    if (circleDistance.y > (rect.height/2 + circle.r)) { return false; }

    if (circleDistance.x <= (rect.width/2)) { return true; } 
    if (circleDistance.y <= (rect.height/2)) { return true; }

    cornerDistance_sq = (circleDistance.x - rect.width/2)^2 +
                         (circleDistance.y - rect.height/2)^2;

    return (cornerDistance_sq <= (circle.r^2));
}

Сперва (первая пара строк) мы сводим вычисления 4-ёх квадрантов к одному. В следующей паре строк проверяем, лежит ли окружность в зеленой области. Если лежит, то пересечения нет. Следующая пара строк проверит, входит ли окружность в оранжевую или серую области рисунка. Если входит, то пересечение обнаружено.

Дальше необходимо проверить, пересекает ли окружность угол прямоугольника (это делают следующие строки кода).

Собственно, этот расчет возвращает 'false' для всех окружностей, центр которых находится в пределах красной области, и 'true' для всех окружностей, центр которых находится в белой области.

В целом, данный код обеспечивает то, что нам и нужно (я привел здесь реализацию кода для 2D, однако в моем коде KD-Дерева я использую 3D вариант).

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

Как было сказано выше, 'веерные' ситуации порождают большое число элементов внутри узла, чем их больше-тем медленнее поиск. Более того, если все элементы равноудалены от данной точки, то поиск будет осуществляться за O(N)(множество точек, которые лежат на поверхности сферы, а ближний ищем к центру сферы). Однако, если данные ситуации убрать, то поиск в среднем будет равносилен спуску по дереву с перебором элементов в нескольких узлах т.е. за O(log(N)). Понятно, что скорость работы поиска зависит от числа посещенных листовых узлов дерева.

Рассмотрим следующие два рисунка:

Суть этих рисунков заключается в том, что если точка к которой мы ищем ближний элемент расположена очень-очень далеко от исходного bounding box-а множества, то сфера с радисом длины minDist(расстояние до ближнего), будет пересекать много больше узлов, чем если бы мы рассматривали эту же сферу, но с центром в точке гораздо ближней к исходному bounding box-у множества (естественно, minDist изменится). В общем, поиск ближних к сильно удаленной точке, осуществляется медленне, чем поиск к точке, расположенной близко к исходному множесту. Мои тесты эту информацию подтвердили.

Результаты и подведение итогов

В качестве итогов хотел бы добавить еще пару слов о своей реализации KD-Дерева и привести полученные результаты. Собственно, дизайн кода разрабатывался так, чтобы его можно было легко адаптировать под любые объекты исходного множества(треугольники, сферы, точки и т.п). Все что нужно-создать класс наследник, с переопределенными виртуальными функциями. Более того, моя реализация также предусматривает передачу специального класса Splitter-а, определяемого пользователем. Этот класс, а точнее его виртуальный метод split, определяют, где конкретно будет проходить секущая плоскость. В своей реализации я предоставляю класс, принимающий решение на основе SAH. Здесь, отмечу, что во многих статьях посвященных ускорению построения KD-Дерева на основе SAH, многие авторы на начальных значениях глубины дерева(вообще, когда количество элементов в узле дерева велико) используют более простые техники поиска секущей плоскости(вроде деления по центру или медиане), а SAH эвристику применяют лишь в тот момент, когда число элементов в узле небольшое.

Моя реализация таких ухищрений не содержит, но позволяет быстро добавить их(нужно лишь расширить новым параметром конструктор KD-Дерева и вызвать функцию-член построения с нужным сплитером, контролируя нужные ограничения). Поиск в дереве выполняю многопоточно. Все вычисления произвожу в числах с двойной точностью(double). Максимальная глубина дерева задается константой(по дефолту-32). Определены некоторые #defines, позволяющие, например, выполнять обход дерева при поиске без рекурсии(с рекурсией, правда быстрее выходт-все узлы это элементы некоторого вектора(то есть расположены рядом по памяти), а значит, хорошо кэшируются). Вместе с кодом я предоставляю тестовые наборы данных(3D модели 'измененного формата OFF' с разным количеством треугольников внутри(от 2 до 3 000 000)). Пользователь может скинуть дамп построенного дерева(формат DXF) и просмотреть его в соответствующей графической программе. Также программа ведет лог(можно включить/выключить) качества построения дерева: сбрасывается глубина дерева, максимальное количество элементов в листовом узле, среднее количесвто элементов в листовых узлах, время работы. Ни в коем случае не утверждаю что полученная реализация идеальна, да что уж там, я и сам знаю места где я промахнулся(например, не передаю аллокатор в параметр шаблона, частое присутствие С-кода(чтение и запись файлов выполняю не с помощью потоков), возможны незамеченные баги и т.п-it's time to fix it). Ну и конечно, дерево сделано и оптимизировано строго для работы в 3D пространстве.

Тестирование кода я проводил на ноутбуке, со следующими характеристиками: Intel Core I7-4750HQ, 4 core(8 threads) 2 GHz, RAM-8gb, Win x64 app на Windows 10. Коэффициенты для вычисления SAH я брал следующие: cT = 1., cI = 1.5. И, если говорть о результатах, то получилось, что на 1, 5млн. треугольников дерево строится менее чем за 1,5 секунды. На 3-ех миллионах за 2,4 секунды. Для 1,5млн. точек и 1.5млн треугольников(точки расположены не очень далеко от исходной модели), поиск выполнился за 0.23 секунды, а если точки удалять от модели-время растет, вплоть до 3 секунд. Для 3-ех миллионов точек(снова, расположены близко к модели) и 3-ех миллионов треугольников поиск занимает около 0,7 секунд. Надеюсь, ничего не перепутал. Напоследок, иллюстрация визуализации построенного KD-Дерева:

Полезные ссылки

[0]: Моя реализация KD-Tree на GitHub
[1]: Поиск расстояния от точки до треугольника
[2]: Определения пересечений окружности и прямоугольника
[3]: Описание KD-Tree на Wikipedia
[4]: Интересная статья посвященная SAH
[5]: Описание KD-Tree на ray-tracing.ru