cp-library
競技プログラミング用 C++ ライブラリ
Kd-tree（最近傍・矩形範囲検索）
data_structure/kd_tree.hpp · GitHub で見る
共通ページ
静的な点集合から最近傍、距離順の上位 k 点、軸に平行な矩形内の点を検索します。次元はコンパイル時に指定し、既定は 2 次元です。
#include "lib/data_structure/kd_tree.hpp"
#include <algorithm>
#include <cassert>
#include <vector>
int main() {
using Tree = KDTree<>;
std::vector<Tree::point_type> points{ {0, 0}, {2, 0}, {2, 0}, {5, 4} };
Tree tree(points);
auto near = tree.nearest({1, 0});
assert(near && near->index == 0 && near->distance_squared == 1);
auto other = tree.nearest(points[1], 1);
assert(other && other->index == 2 && other->distance_squared == 0);
auto closest = tree.k_nearest({1, 0}, 3);
assert(closest.size() == 3);
assert(closest[0].index == 0 && closest[1].index == 1 && closest[2].index == 2);
auto ids = tree.range_report({0, -1}, {3, 1});
std::sort(ids.begin(), ids.end());
assert((ids == std::vector<int>{0, 1, 2}));
assert(tree.range_count({0, -1}, {2, 1}) == 1);
int count = 0;
tree.range_visit({0, -1}, {3, 1}, [&](int id) {
assert(tree.point(id)[1] == 0);
++count;
});
assert(count == 3);
}
KDTree<T = std::int64_t, Dimension = 2, Distance = ...> を使います。T は座標型、Dimension は正の int、Distance は距離二乗の型です。Distance の既定値は整数座標では __int128_t、浮動小数点座標では T です。
point_type は std::array<T, Dimension>、distance_type は Distance です。neighbor_type は int index と Distance distance_squared を持ち、等値比較もできます。距離はユークリッド距離の二乗で、平方根は取りません。
点の番号は構築時の入力順に [0,size()) です。同じ座標の点も別々の番号で保持します。内部で点を並べ替えても番号は変わりません。入力はコピーして保持するため、構築後に元の配列を変更・破棄できます。
| 操作 | 内容・返り値 |
|---|---|
KDTree() |
空の木。葉の上限点数は 16 |
KDTree(points, leaf_size = 16) |
std::span<const point_type> から構築。vector<point_type> や array<point_type,N> も渡せる |
build(points, leaf_size = 16) |
点集合を置き換えて構築し直す。番号も入力順に振り直す。返り値は void
|
size() / empty()
|
点数を int／空かを bool で返す |
leaf_size() |
葉に保持する点数の上限を int で返す |
point(index) |
入力番号 index の座標を const point_type& で返す。番号は [0,size())
|
nearest(query, exclude = -1) |
最近傍を std::optional<neighbor_type> で返す。対象点がなければ std::nullopt
|
k_nearest(query, k, exclude = -1) |
距離二乗と番号の昇順に、最大 k 点を std::vector<neighbor_type> で返す |
range_count(lower, upper) |
各軸の半開区間 [lower[d],upper[d]) に入る点数を int で返す |
range_report(lower, upper) |
同じ半開矩形に入る点の番号を std::vector<int> で返す。順序は未規定 |
range_visit(lower, upper, visit) |
同じ半開矩形の各点について visit(int index) を 1 回ずつ呼ぶ。返り値は void、順序は未規定 |
点数は INT_MAX 以下、leaf_size は正の int とします。空の入力も構築できます。build で leaf_size を省略すると、前回の設定によらず 16 になります。葉を大きくすると木のノード数が減り、葉で直接比較する点数が増えます。
query・lower・upper は point_type です。nearest と k_nearest は、距離二乗が同じなら入力番号の小さい点を優先します。exclude は -1 または [0,size()) の int で、指定した番号の 1 点だけを除外します。同じ座標にある他の点は検索対象です。k は非負の int で、対象点数より大きい場合は全対象点を返します。k == 0 または対象点がない場合は空の列です。
矩形は全軸で下端を含み、上端を含みません。どれかの軸で lower[d] >= upper[d] なら空として扱います。該当する点がなければ range_count は 0、range_report は空の列、range_visit は呼び出しなしになります。range_visit のコールバックでは木を変更しないでください。番号順が必要なら、取得した番号列をソートします。
検索と取得は const な木から呼び出せます。点の追加・削除・座標変更は build で点集合全体を置き換えます。コピーした木は独立して使えます。point が返す参照は、再構築・代入・破棄で無効になります。
座標には bool を除く 64 bit 以下の符号付き・符号なし整数、または組み込み浮動小数点型を使います。整数座標では Distance も bool を除く 128 bit 以下の整数型、浮動小数点座標では T 以上の精度を持つ浮動小数点型を指定します。
整数では、問い合わせ点と各登録点の座標差の二乗、およびその全軸の和が Distance に収まる必要があります。減算は拡張した型で行いますが、距離二乗のオーバーフローは検出しません。既定の 128 bit 距離は 2 次元の 32 bit 整数座標の全範囲を扱えます。2 次元で登録点・問い合わせ点の各座標の絶対値が \(10^{18}\) 以下なら、64 bit 座標と既定の距離型で収まります。距離の範囲が既知なら、例えば KDTree<std::int64_t, 2, std::int64_t> として距離型を狭くできます。この場合、2 次元で各座標の絶対値が \(10^9\) 以下なら距離二乗は収まります。
浮動小数点では、登録点と最近傍の問い合わせ点の座標は有限値にします。座標差・距離二乗などの中間値も有限となる範囲を使ってください。矩形の上下端には正負の無限大も使えますが、NaN は扱いません。比較に EPS は使わず、浮動小数点の丸めに従います。
3 次元の浮動小数点座標でも同じ API を使えます。
#include "lib/data_structure/kd_tree.hpp"
#include <cassert>
#include <limits>
#include <vector>
int main() {
using Tree = KDTree<double, 3>;
std::vector<Tree::point_type> points{ {1.0, 2.0, 3.0}, {-1.0, 0.0, 0.0} };
Tree tree(points, 8);
auto near = tree.nearest({0.0, 0.0, 0.0});
assert(near && near->index == 1 && near->distance_squared == 1.0);
double inf = std::numeric_limits<double>::infinity();
assert(tree.range_count({-inf, -inf, -inf}, {inf, inf, inf}) == 2);
}
座標の広がりが最も大きい軸で中央値に分割し、点数が釣り合う木を作ります。点とノードを連続した配列に保持し、葉では点を直接比較します。部分木の外接矩形と最小の入力番号を使って、より良い答えを含まない部分木を枝刈りします。矩形全体が検索範囲に入る部分木は、個数をまとめて数えられます。
参考: Jon Louis Bentley, Multidimensional binary search trees used for associative searching, Communications of the ACM, 1975.
点数を \(n\)、次元を \(D\)、返す近傍点数を \(k\)、矩形に入る点数を \(z\) とします。std::nth_element による構築・build は平均 \(O(Dn\log n)\) 時間で、保持空間は \(O(Dn)\) です。葉の上限点数を超えない場合の構築は \(O(Dn)\) です。
最近傍と矩形の個数取得は最悪 \(O(Dn)\)、矩形の列挙・コールバックは最悪 \(O(Dn+z)\) 時間です。k_nearest は \(k\ge2\) で最悪 \(O(Dn+n\log k)\)、\(k=1\) では最近傍と同じです。次元を固定すると、空入力の検索と k == 0 は定数時間です。軸を広がりで選ぶため、矩形検索に \(O(n^{1-1/D}+z)\) の最悪計算量は保証しません。
検索の作業領域は再帰スタックに \(O(\log n)\)、k_nearest では候補ヒープにさらに \(O(k)\) を使います。返す列の領域は別途 \(O(k)\) または \(O(z)\) です。size・empty・leaf_size・point は \(O(1)\) 時間です。空集合・点が 1 個の場合の管理操作は定数時間です。座標の演算とコールバック 1 回を定数時間とします。
依存ライブラリ
検証コード
コード
#ifndef KD_TREE_HPP
#define KD_TREE_HPP
#include "../components/d_ary_heap.hpp"
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <span>
#include <type_traits>
#include <utility>
#include <vector>
template<class T = std::int64_t, int Dimension = 2,
class Distance = std::conditional_t<std::is_integral_v<T>, __int128_t, T>>
class KDTree {
static_assert(Dimension > 0);
static_assert(std::is_arithmetic_v<T> && !std::is_same_v<T, bool>);
static_assert(!std::is_integral_v<T> || sizeof(T) <= 8);
static_assert(std::numeric_limits<Distance>::is_bounded && !std::is_same_v<Distance, bool>);
static_assert(std::numeric_limits<Distance>::is_integer == std::is_integral_v<T>);
static_assert(std::numeric_limits<Distance>::is_integer ? sizeof(Distance) <= 16 :
(std::is_floating_point_v<Distance> && std::numeric_limits<Distance>::digits >= std::numeric_limits<T>::digits));
public:
using point_type = std::array<T, Dimension>;
using distance_type = Distance;
struct neighbor_type {
int index;
Distance distance_squared;
bool operator==(const neighbor_type&) const = default;
};
private:
using Wide = std::conditional_t<std::is_integral_v<T>, std::conditional_t<(sizeof(T) < 8), std::int64_t, __int128_t>, Distance>;
struct Entry { point_type point; int index; };
struct Node { point_type lower, upper; int begin, end, minimum; std::uint32_t right = 0; };
struct Less {
bool operator()(const neighbor_type& a, const neighbor_type& b) const {
return a.distance_squared < b.distance_squared || (a.distance_squared == b.distance_squared && a.index < b.index);
}
};
using Heap = DaryHeap<neighbor_type, Less, 4>;
std::vector<Entry> points_;
std::vector<int> positions_;
std::vector<Node> nodes_;
int leaf_size_ = 16;
static void check_point([[maybe_unused]] const point_type& point) {
if constexpr (std::is_floating_point_v<T>)
for ([[maybe_unused]] T value : point) assert(std::isfinite(value));
}
void check_exclude([[maybe_unused]] int exclude) const { assert(-1 <= exclude && exclude < size()); }
static Distance square_difference(T a, T b) {
Wide difference = Wide(a) - Wide(b);
Distance value = Distance(difference);
if constexpr (std::numeric_limits<Distance>::is_integer && sizeof(Distance) < sizeof(unsigned))
return Distance(unsigned(value) * unsigned(value));
return value * value;
}
static Distance distance(const point_type& a, const point_type& b) {
Distance result = 0;
for (int d = 0; d < Dimension; ++d) result += square_difference(a[d], b[d]);
return result;
}
static Distance lower_bound(const Node& node, const point_type& query) {
Distance result = 0;
for (int d = 0; d < Dimension; ++d) {
if (query[d] < node.lower[d]) result += square_difference(query[d], node.lower[d]);
else if (node.upper[d] < query[d]) result += square_difference(query[d], node.upper[d]);
}
return result;
}
std::uint32_t construct(int begin, int end) {
std::uint32_t id = std::uint32_t(nodes_.size());
Node node{points_[begin].point, points_[begin].point, begin, end, points_[begin].index};
for (int i = begin + 1; i < end; ++i) {
node.minimum = std::min(node.minimum, points_[i].index);
for (int d = 0; d < Dimension; ++d) {
node.lower[d] = std::min(node.lower[d], points_[i].point[d]);
node.upper[d] = std::max(node.upper[d], points_[i].point[d]);
}
}
nodes_.push_back(node);
if (end - begin <= leaf_size_) return id;
int axis = 0;
for (int d = 1; d < Dimension; ++d)
if (Wide(node.upper[axis]) - Wide(node.lower[axis]) < Wide(node.upper[d]) - Wide(node.lower[d])) axis = d;
int middle = begin + (end - begin) / 2;
std::nth_element(points_.begin() + begin, points_.begin() + middle, points_.begin() + end,
[axis](const Entry& a, const Entry& b) {
return a.point[axis] < b.point[axis] || (a.point[axis] == b.point[axis] && a.index < b.index);
});
construct(begin, middle);
nodes_[id].right = construct(middle, end);
return id;
}
bool promising(Distance bound, std::uint32_t id, const neighbor_type& best) const {
return Less{}({nodes_[id].minimum, bound}, best);
}
void nearest_node(std::uint32_t id, const point_type& query, int exclude, neighbor_type& best) const {
const auto& node = nodes_[id];
if (!node.right) {
for (int i = node.begin; i < node.end; ++i) {
const auto& point = points_[i];
if (point.index == exclude) continue;
neighbor_type candidate{point.index, distance(query, point.point)};
if (Less{}(candidate, best)) best = candidate;
}
return;
}
std::uint32_t first = id + 1, second = node.right;
Distance a = lower_bound(nodes_[first], query), b = lower_bound(nodes_[second], query);
if (Less{}({nodes_[second].minimum, b}, {nodes_[first].minimum, a})) std::swap(first, second), std::swap(a, b);
if (promising(a, first, best)) nearest_node(first, query, exclude, best);
if (promising(b, second, best)) nearest_node(second, query, exclude, best);
}
void nearest_k_node(std::uint32_t id, const point_type& query, int k, int exclude, Heap& heap) const {
const auto& node = nodes_[id];
if (!node.right) {
for (int i = node.begin; i < node.end; ++i) {
const auto& point = points_[i];
if (point.index == exclude) continue;
neighbor_type candidate{point.index, distance(query, point.point)};
if (int(heap.size()) < k) heap.push(candidate);
else if (Less{}(candidate, heap.top())) heap.replace_top(candidate);
}
return;
}
std::uint32_t first = id + 1, second = node.right;
Distance a = lower_bound(nodes_[first], query), b = lower_bound(nodes_[second], query);
if (Less{}({nodes_[second].minimum, b}, {nodes_[first].minimum, a})) std::swap(first, second), std::swap(a, b);
if (int(heap.size()) < k || promising(a, first, heap.top())) nearest_k_node(first, query, k, exclude, heap);
if (int(heap.size()) < k || promising(b, second, heap.top())) nearest_k_node(second, query, k, exclude, heap);
}
static bool nonempty_box(const point_type& lower, const point_type& upper) {
for (int d = 0; d < Dimension; ++d) if (!(lower[d] < upper[d])) return false;
return true;
}
template<bool Report, class F>
int range_node(std::uint32_t id, const point_type& lower, const point_type& upper, F& visit) const {
const auto& node = nodes_[id];
bool inside = true;
for (int d = 0; d < Dimension; ++d) {
if (node.upper[d] < lower[d] || !(node.lower[d] < upper[d])) return 0;
inside &= !(node.lower[d] < lower[d]) && node.upper[d] < upper[d];
}
if (inside) {
if constexpr (Report) for (int i = node.begin; i < node.end; ++i) visit(points_[i].index);
return node.end - node.begin;
}
if (node.right) return range_node<Report>(id + 1, lower, upper, visit) + range_node<Report>(node.right, lower, upper, visit);
int count = 0;
for (int i = node.begin; i < node.end; ++i) {
bool found = true;
for (int d = 0; d < Dimension; ++d)
found &= !(points_[i].point[d] < lower[d]) && points_[i].point[d] < upper[d];
if (found) {
if constexpr (Report) visit(points_[i].index);
++count;
}
}
return count;
}
public:
KDTree() = default;
explicit KDTree(std::span<const point_type> points, int leaf_size = 16) { build(points, leaf_size); }
void build(std::span<const point_type> points, int leaf_size = 16) {
assert(points.size() <= std::size_t(std::numeric_limits<int>::max()) && leaf_size > 0);
KDTree next;
next.leaf_size_ = leaf_size;
next.points_.reserve(points.size());
for (int i = 0; i < int(points.size()); ++i) {
check_point(points[i]);
next.points_.push_back({points[i], i});
}
next.positions_.resize(points.size());
if (!points.empty()) {
next.nodes_.reserve(std::min(points.size() * 2, (points.size() / std::size_t(leaf_size) + 1) * 4));
next.construct(0, int(points.size()));
for (int i = 0; i < int(points.size()); ++i) next.positions_[next.points_[i].index] = i;
}
*this = std::move(next);
}
int size() const { return int(points_.size()); }
bool empty() const { return points_.empty(); }
int leaf_size() const { return leaf_size_; }
const point_type& point(int index) const {
assert(0 <= index && index < size());
return points_[positions_[index]].point;
}
std::optional<neighbor_type> nearest(const point_type& query, int exclude = -1) const {
check_point(query);
check_exclude(exclude);
if (empty() || (size() == 1 && exclude == 0)) return std::nullopt;
neighbor_type best{std::numeric_limits<int>::max(), std::numeric_limits<Distance>::max()};
nearest_node(0, query, exclude, best);
return best;
}
std::vector<neighbor_type> k_nearest(const point_type& query, int k, int exclude = -1) const {
assert(k >= 0);
check_point(query);
check_exclude(exclude);
k = std::min(k, size() - (exclude != -1));
if (!k) return {};
if (k == 1) return {*nearest(query, exclude)};
Heap heap;
heap.reserve(k);
nearest_k_node(0, query, k, exclude, heap);
auto result = heap.values();
std::sort(result.begin(), result.end(), Less{});
return result;
}
int range_count(const point_type& lower, const point_type& upper) const {
if (empty() || !nonempty_box(lower, upper)) return 0;
auto visit = [](int) {};
return range_node<false>(0, lower, upper, visit);
}
template<class F> void range_visit(const point_type& lower, const point_type& upper, F&& visit) const {
if (!empty() && nonempty_box(lower, upper)) range_node<true>(0, lower, upper, visit);
}
std::vector<int> range_report(const point_type& lower, const point_type& upper) const {
std::vector<int> result;
range_visit(lower, upper, [&](int index) { result.push_back(index); });
return result;
}
};
#endif
#line 1 "data_structure/kd_tree.hpp"
#line 1 "components/d_ary_heap.hpp"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <functional>
#include <utility>
#include <vector>
template <class T, class Compare = std::less<T>, std::size_t D = 2>
class DaryHeap {
static_assert(D >= 2);
std::vector<T> data;
[[no_unique_address]] Compare compare{};
void up(std::size_t i, T value) {
while (i) {
std::size_t p = (i - 1) / D;
if (!compare(data[p], value)) break;
data[i] = std::move(data[p]);
i = p;
}
data[i] = std::move(value);
}
void down(std::size_t i, T value) {
std::size_t n = data.size();
while (i < n / D + (n % D > 1)) {
std::size_t best = i * D + 1, end = n;
if constexpr (D <= std::size_t(-1) / 2) end = std::min(best + D, n);
for (std::size_t j = best + 1; j < end; ++j)
if (compare(data[best], data[j])) best = j;
if (!compare(value, data[best])) break;
data[i] = std::move(data[best]);
i = best;
}
data[i] = std::move(value);
}
public:
DaryHeap() = default;
explicit DaryHeap(Compare compare) : compare(std::move(compare)) {}
explicit DaryHeap(std::vector<T> values, Compare compare = {}) : data(std::move(values)), compare(std::move(compare)) {
for (std::size_t i = data.size() / D + (data.size() % D > 1); i--;) down(i, std::move(data[i]));
}
bool empty() const { return data.empty(); }
std::size_t size() const { return data.size(); }
void reserve(std::size_t n) { data.reserve(n); }
void clear() { data.clear(); }
const std::vector<T> &values() const { return data; }
decltype(auto) top() const {
assert(!empty());
return data.front();
}
void push(T value) {
data.push_back(std::move(value));
up(data.size() - 1, std::move(data.back()));
}
T replace_top(T value) {
assert(!empty());
T result = std::move(data.front());
down(0, std::move(value));
return result;
}
T push_pop(T value) {
if (empty() || !compare(value, data.front())) return value;
return replace_top(std::move(value));
}
T pop_value() {
assert(!empty());
T result = std::move(data.front());
if (data.size() == 1) data.pop_back();
else {
T value = std::move(data.back());
data.pop_back();
down(0, std::move(value));
}
return result;
}
void pop() { pop_value(); }
void swap(DaryHeap &other) noexcept(noexcept(std::swap(compare, other.compare))) {
data.swap(other.data);
std::swap(compare, other.compare);
}
};
#line 6 "data_structure/kd_tree.hpp"
#include <array>
#line 8 "data_structure/kd_tree.hpp"
#include <cmath>
#line 10 "data_structure/kd_tree.hpp"
#include <cstdint>
#include <limits>
#include <optional>
#include <span>
#include <type_traits>
#line 17 "data_structure/kd_tree.hpp"
template<class T = std::int64_t, int Dimension = 2,
class Distance = std::conditional_t<std::is_integral_v<T>, __int128_t, T>>
class KDTree {
static_assert(Dimension > 0);
static_assert(std::is_arithmetic_v<T> && !std::is_same_v<T, bool>);
static_assert(!std::is_integral_v<T> || sizeof(T) <= 8);
static_assert(std::numeric_limits<Distance>::is_bounded && !std::is_same_v<Distance, bool>);
static_assert(std::numeric_limits<Distance>::is_integer == std::is_integral_v<T>);
static_assert(std::numeric_limits<Distance>::is_integer ? sizeof(Distance) <= 16 :
(std::is_floating_point_v<Distance> && std::numeric_limits<Distance>::digits >= std::numeric_limits<T>::digits));
public:
using point_type = std::array<T, Dimension>;
using distance_type = Distance;
struct neighbor_type {
int index;
Distance distance_squared;
bool operator==(const neighbor_type&) const = default;
};
private:
using Wide = std::conditional_t<std::is_integral_v<T>, std::conditional_t<(sizeof(T) < 8), std::int64_t, __int128_t>, Distance>;
struct Entry { point_type point; int index; };
struct Node { point_type lower, upper; int begin, end, minimum; std::uint32_t right = 0; };
struct Less {
bool operator()(const neighbor_type& a, const neighbor_type& b) const {
return a.distance_squared < b.distance_squared || (a.distance_squared == b.distance_squared && a.index < b.index);
}
};
using Heap = DaryHeap<neighbor_type, Less, 4>;
std::vector<Entry> points_;
std::vector<int> positions_;
std::vector<Node> nodes_;
int leaf_size_ = 16;
static void check_point([[maybe_unused]] const point_type& point) {
if constexpr (std::is_floating_point_v<T>)
for ([[maybe_unused]] T value : point) assert(std::isfinite(value));
}
void check_exclude([[maybe_unused]] int exclude) const { assert(-1 <= exclude && exclude < size()); }
static Distance square_difference(T a, T b) {
Wide difference = Wide(a) - Wide(b);
Distance value = Distance(difference);
if constexpr (std::numeric_limits<Distance>::is_integer && sizeof(Distance) < sizeof(unsigned))
return Distance(unsigned(value) * unsigned(value));
return value * value;
}
static Distance distance(const point_type& a, const point_type& b) {
Distance result = 0;
for (int d = 0; d < Dimension; ++d) result += square_difference(a[d], b[d]);
return result;
}
static Distance lower_bound(const Node& node, const point_type& query) {
Distance result = 0;
for (int d = 0; d < Dimension; ++d) {
if (query[d] < node.lower[d]) result += square_difference(query[d], node.lower[d]);
else if (node.upper[d] < query[d]) result += square_difference(query[d], node.upper[d]);
}
return result;
}
std::uint32_t construct(int begin, int end) {
std::uint32_t id = std::uint32_t(nodes_.size());
Node node{points_[begin].point, points_[begin].point, begin, end, points_[begin].index};
for (int i = begin + 1; i < end; ++i) {
node.minimum = std::min(node.minimum, points_[i].index);
for (int d = 0; d < Dimension; ++d) {
node.lower[d] = std::min(node.lower[d], points_[i].point[d]);
node.upper[d] = std::max(node.upper[d], points_[i].point[d]);
}
}
nodes_.push_back(node);
if (end - begin <= leaf_size_) return id;
int axis = 0;
for (int d = 1; d < Dimension; ++d)
if (Wide(node.upper[axis]) - Wide(node.lower[axis]) < Wide(node.upper[d]) - Wide(node.lower[d])) axis = d;
int middle = begin + (end - begin) / 2;
std::nth_element(points_.begin() + begin, points_.begin() + middle, points_.begin() + end,
[axis](const Entry& a, const Entry& b) {
return a.point[axis] < b.point[axis] || (a.point[axis] == b.point[axis] && a.index < b.index);
});
construct(begin, middle);
nodes_[id].right = construct(middle, end);
return id;
}
bool promising(Distance bound, std::uint32_t id, const neighbor_type& best) const {
return Less{}({nodes_[id].minimum, bound}, best);
}
void nearest_node(std::uint32_t id, const point_type& query, int exclude, neighbor_type& best) const {
const auto& node = nodes_[id];
if (!node.right) {
for (int i = node.begin; i < node.end; ++i) {
const auto& point = points_[i];
if (point.index == exclude) continue;
neighbor_type candidate{point.index, distance(query, point.point)};
if (Less{}(candidate, best)) best = candidate;
}
return;
}
std::uint32_t first = id + 1, second = node.right;
Distance a = lower_bound(nodes_[first], query), b = lower_bound(nodes_[second], query);
if (Less{}({nodes_[second].minimum, b}, {nodes_[first].minimum, a})) std::swap(first, second), std::swap(a, b);
if (promising(a, first, best)) nearest_node(first, query, exclude, best);
if (promising(b, second, best)) nearest_node(second, query, exclude, best);
}
void nearest_k_node(std::uint32_t id, const point_type& query, int k, int exclude, Heap& heap) const {
const auto& node = nodes_[id];
if (!node.right) {
for (int i = node.begin; i < node.end; ++i) {
const auto& point = points_[i];
if (point.index == exclude) continue;
neighbor_type candidate{point.index, distance(query, point.point)};
if (int(heap.size()) < k) heap.push(candidate);
else if (Less{}(candidate, heap.top())) heap.replace_top(candidate);
}
return;
}
std::uint32_t first = id + 1, second = node.right;
Distance a = lower_bound(nodes_[first], query), b = lower_bound(nodes_[second], query);
if (Less{}({nodes_[second].minimum, b}, {nodes_[first].minimum, a})) std::swap(first, second), std::swap(a, b);
if (int(heap.size()) < k || promising(a, first, heap.top())) nearest_k_node(first, query, k, exclude, heap);
if (int(heap.size()) < k || promising(b, second, heap.top())) nearest_k_node(second, query, k, exclude, heap);
}
static bool nonempty_box(const point_type& lower, const point_type& upper) {
for (int d = 0; d < Dimension; ++d) if (!(lower[d] < upper[d])) return false;
return true;
}
template<bool Report, class F>
int range_node(std::uint32_t id, const point_type& lower, const point_type& upper, F& visit) const {
const auto& node = nodes_[id];
bool inside = true;
for (int d = 0; d < Dimension; ++d) {
if (node.upper[d] < lower[d] || !(node.lower[d] < upper[d])) return 0;
inside &= !(node.lower[d] < lower[d]) && node.upper[d] < upper[d];
}
if (inside) {
if constexpr (Report) for (int i = node.begin; i < node.end; ++i) visit(points_[i].index);
return node.end - node.begin;
}
if (node.right) return range_node<Report>(id + 1, lower, upper, visit) + range_node<Report>(node.right, lower, upper, visit);
int count = 0;
for (int i = node.begin; i < node.end; ++i) {
bool found = true;
for (int d = 0; d < Dimension; ++d)
found &= !(points_[i].point[d] < lower[d]) && points_[i].point[d] < upper[d];
if (found) {
if constexpr (Report) visit(points_[i].index);
++count;
}
}
return count;
}
public:
KDTree() = default;
explicit KDTree(std::span<const point_type> points, int leaf_size = 16) { build(points, leaf_size); }
void build(std::span<const point_type> points, int leaf_size = 16) {
assert(points.size() <= std::size_t(std::numeric_limits<int>::max()) && leaf_size > 0);
KDTree next;
next.leaf_size_ = leaf_size;
next.points_.reserve(points.size());
for (int i = 0; i < int(points.size()); ++i) {
check_point(points[i]);
next.points_.push_back({points[i], i});
}
next.positions_.resize(points.size());
if (!points.empty()) {
next.nodes_.reserve(std::min(points.size() * 2, (points.size() / std::size_t(leaf_size) + 1) * 4));
next.construct(0, int(points.size()));
for (int i = 0; i < int(points.size()); ++i) next.positions_[next.points_[i].index] = i;
}
*this = std::move(next);
}
int size() const { return int(points_.size()); }
bool empty() const { return points_.empty(); }
int leaf_size() const { return leaf_size_; }
const point_type& point(int index) const {
assert(0 <= index && index < size());
return points_[positions_[index]].point;
}
std::optional<neighbor_type> nearest(const point_type& query, int exclude = -1) const {
check_point(query);
check_exclude(exclude);
if (empty() || (size() == 1 && exclude == 0)) return std::nullopt;
neighbor_type best{std::numeric_limits<int>::max(), std::numeric_limits<Distance>::max()};
nearest_node(0, query, exclude, best);
return best;
}
std::vector<neighbor_type> k_nearest(const point_type& query, int k, int exclude = -1) const {
assert(k >= 0);
check_point(query);
check_exclude(exclude);
k = std::min(k, size() - (exclude != -1));
if (!k) return {};
if (k == 1) return {*nearest(query, exclude)};
Heap heap;
heap.reserve(k);
nearest_k_node(0, query, k, exclude, heap);
auto result = heap.values();
std::sort(result.begin(), result.end(), Less{});
return result;
}
int range_count(const point_type& lower, const point_type& upper) const {
if (empty() || !nonempty_box(lower, upper)) return 0;
auto visit = [](int) {};
return range_node<false>(0, lower, upper, visit);
}
template<class F> void range_visit(const point_type& lower, const point_type& upper, F&& visit) const {
if (!empty() && nonempty_box(lower, upper)) range_node<true>(0, lower, upper, visit);
}
std::vector<int> range_report(const point_type& lower, const point_type& upper) const {
std::vector<int> result;
range_visit(lower, upper, [&](int index) { result.push_back(index); });
return result;
}
};