kd_tree.hpp¶
Static two-dimensional KD-tree for rectangle reporting and nearest neighbor queries.
维护低维点集的空间划分，用于矩形范围搜索或最近邻类剪枝；适合维数很小的几何查询。
Implementation¶
#ifndef NOYA_KD_TREE_HPP
#define NOYA_KD_TREE_HPP 1
/// @complexity Time: Expected O(n log n) build; query time is output/pruning dependent, worst O(n).
/// Space: O(n).
#include "noya/geometry_base.hpp"
#include <algorithm>
#include <cassert>
#include <numeric>
#include <type_traits>
#include <utility>
#include <vector>
namespace noya {
/// @brief Static two-dimensional KD-tree for rectangle reporting and nearest
/// neighbor queries.
template <class T, class Distance = std::conditional_t<std::is_integral_v<T>,
__int128_t, long double>>
struct kd_tree {
using point_type = point<T>;
using distance_type = Distance;
struct node {
int point_index = -1;
int left = -1;
int right = -1;
T min_x{};
T max_x{};
T min_y{};
T max_y{};
};
std::vector<point_type> points;
std::vector<node> nodes;
int root = -1;
kd_tree() = default;
explicit kd_tree(const std::vector<point_type> &input) { build(input); }
/// @brief Rebuild the balanced tree in O(n log n) expected time.
void build(const std::vector<point_type> &input) {
points = input;
nodes.clear();
root = -1;
std::vector<int> order(points.size());
std::iota(order.begin(), order.end(), 0);
if (!order.empty()) {
nodes.reserve(order.size());
root = build_range(order, 0, int(order.size()), 0);
}
}
/// @brief Return the number of stored points.
int size() const { return int(points.size()); }
/// @brief Return point indices in [left, right) x [bottom, top), stopping
/// after max_count results when max_count is nonnegative.
std::vector<int> range_query(const T &left, const T &right, const T &bottom,
const T &top, int max_count = -1) const {
assert(!(right < left) && !(top < bottom));
int limit = max_count < 0 ? size() : max_count;
std::vector<int> result;
result.reserve(std::min(size(), limit));
range_query_at(root, left, right, bottom, top, limit, result);
return result;
}
/// @brief Return (point index, squared distance) for a nearest point; the
/// index is -1 when the tree is empty.
std::pair<int, distance_type> nearest(const point_type &query) const {
int best_index = -1;
distance_type best_distance{};
nearest_at(root, query, best_index, best_distance);
return {best_index, best_distance};
}
private:
static distance_type square_difference(const T &first, const T &second) {
distance_type difference =
static_cast<distance_type>(first) - static_cast<distance_type>(second);
return difference * difference;
}
distance_type point_distance(int index, const point_type &query) const {
return square_difference(points[index].x, query.x) +
square_difference(points[index].y, query.y);
}
distance_type box_distance(const node ¤t,
const point_type &query) const {
distance_type result{};
if (query.x < current.min_x) {
result += square_difference(query.x, current.min_x);
} else if (current.max_x < query.x) {
result += square_difference(query.x, current.max_x);
}
if (query.y < current.min_y) {
result += square_difference(query.y, current.min_y);
} else if (current.max_y < query.y) {
result += square_difference(query.y, current.max_y);
}
return result;
}
int build_range(std::vector<int> &order, int left, int right, int depth) {
int middle = (left + right) / 2;
bool split_x = (depth & 1) == 0;
auto compare = [&](int first, int second) {
if (split_x) {
if (points[first].x != points[second].x) {
return points[first].x < points[second].x;
}
if (points[first].y != points[second].y) {
return points[first].y < points[second].y;
}
} else {
if (points[first].y != points[second].y) {
return points[first].y < points[second].y;
}
if (points[first].x != points[second].x) {
return points[first].x < points[second].x;
}
}
return first < second;
};
std::nth_element(order.begin() + left, order.begin() + middle,
order.begin() + right, compare);
int id = int(nodes.size());
int point_index = order[middle];
nodes.push_back({point_index, -1, -1, points[point_index].x,
points[point_index].x, points[point_index].y,
points[point_index].y});
if (left < middle) {
nodes[id].left = build_range(order, left, middle, depth + 1);
extend_box(nodes[id], nodes[nodes[id].left]);
}
if (middle + 1 < right) {
nodes[id].right = build_range(order, middle + 1, right, depth + 1);
extend_box(nodes[id], nodes[nodes[id].right]);
}
return id;
}
static void extend_box(node &target, const node &source) {
target.min_x = std::min(target.min_x, source.min_x);
target.max_x = std::max(target.max_x, source.max_x);
target.min_y = std::min(target.min_y, source.min_y);
target.max_y = std::max(target.max_y, source.max_y);
}
static bool box_disjoint(const node ¤t, const T &left, const T &right,
const T &bottom, const T &top) {
return current.max_x < left || !(current.min_x < right) ||
current.max_y < bottom || !(current.min_y < top);
}
static bool box_contained(const node ¤t, const T &left, const T &right,
const T &bottom, const T &top) {
return !(current.min_x < left) && current.max_x < right &&
!(current.min_y < bottom) && current.max_y < top;
}
void collect_subtree(int id, int limit, std::vector<int> &result) const {
if (id == -1 || int(result.size()) == limit) {
return;
}
result.push_back(nodes[id].point_index);
collect_subtree(nodes[id].left, limit, result);
collect_subtree(nodes[id].right, limit, result);
}
void range_query_at(int id, const T &left, const T &right, const T &bottom,
const T &top, int limit, std::vector<int> &result) const {
if (id == -1 || int(result.size()) == limit ||
box_disjoint(nodes[id], left, right, bottom, top)) {
return;
}
if (box_contained(nodes[id], left, right, bottom, top)) {
collect_subtree(id, limit, result);
return;
}
int index = nodes[id].point_index;
const point_type &value = points[index];
if (!(value.x < left) && value.x < right && !(value.y < bottom) &&
value.y < top) {
result.push_back(index);
}
range_query_at(nodes[id].left, left, right, bottom, top, limit, result);
range_query_at(nodes[id].right, left, right, bottom, top, limit, result);
}
void nearest_at(int id, const point_type &query, int &best_index,
distance_type &best_distance) const {
if (id == -1) {
return;
}
int index = nodes[id].point_index;
distance_type distance = point_distance(index, query);
if (best_index == -1 || distance < best_distance ||
(distance == best_distance && index < best_index)) {
best_index = index;
best_distance = distance;
}
int first = nodes[id].left;
int second = nodes[id].right;
distance_type first_bound =
first == -1 ? distance_type{} : box_distance(nodes[first], query);
distance_type second_bound =
second == -1 ? distance_type{} : box_distance(nodes[second], query);
if (first == -1 || (second != -1 && second_bound < first_bound)) {
std::swap(first, second);
std::swap(first_bound, second_bound);
}
if (first != -1 && first_bound <= best_distance) {
nearest_at(first, query, best_index, best_distance);
}
if (second != -1 && second_bound <= best_distance) {
nearest_at(second, query, best_index, best_distance);
}
}
};
} // namespace noya
#endif // NOYA_KD_TREE_HPP
#include <algorithm>
#include <cassert>
#include <cmath>
#include <numeric>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
/// @complexity Time: Expected O(n log n) build; query time is output/pruning dependent, worst O(n).
/// Space: O(n).
/// @complexity Time: O(1) per primitive; O(n log n) for convex hull.
/// Space: O(1) per primitive and O(n) for hull construction.
namespace noya {
/// @brief Two-dimensional point with vector arithmetic and lexicographic order.
template <class T> struct point {
T x{};
T y{};
point() = default;
point(T x_, T y_) : x(x_), y(y_) {}
point &operator+=(const point &other) {
x += other.x;
y += other.y;
return *this;
}
point &operator-=(const point &other) {
x -= other.x;
y -= other.y;
return *this;
}
point &operator*=(const T &scale) {
x *= scale;
y *= scale;
return *this;
}
point &operator/=(const T &scale) {
x /= scale;
y /= scale;
return *this;
}
friend point operator+(point left, const point &right) {
return left += right;
}
friend point operator-(point left, const point &right) {
return left -= right;
}
friend point operator*(point value, const T &scale) { return value *= scale; }
friend point operator*(const T &scale, point value) { return value *= scale; }
friend point operator/(point value, const T &scale) { return value /= scale; }
friend bool operator==(const point &, const point &) = default;
friend bool operator<(const point &left, const point &right) {
return left.x < right.x || (left.x == right.x && left.y < right.y);
}
};
/// @brief Circle represented by a center and a nonnegative radius.
template <class Real> struct circle {
point<Real> center;
Real radius{};
};
/// @brief Return the dot product of two vectors.
template <class T> T dot(const point<T> &a, const point<T> &b) {
return a.x * b.x + a.y * b.y;
}
/// @brief Return the signed cross product of two vectors.
template <class T> T cross(const point<T> &a, const point<T> &b) {
return a.x * b.y - a.y * b.x;
}
/// @brief Return cross(a - origin, b - origin).
template <class T>
T cross(const point<T> &origin, const point<T> &a, const point<T> &b) {
return cross(a - origin, b - origin);
}
/// @brief Return the squared Euclidean norm.
template <class T> T norm2(const point<T> &value) { return dot(value, value); }
/// @brief Compare a value with zero using an optional absolute tolerance.
template <class T> int sign(const T &value, const T &epsilon = T{}) {
return (value > epsilon) - (value < -epsilon);
}
/// @brief Return -1, 0, or 1 for a clockwise, collinear, or counter-clockwise
/// turn.
template <class T>
int orientation(const point<T> &a, const point<T> &b, const point<T> &c,
const T &epsilon = T{}) {
return sign(cross(a, b, c), epsilon);
}
/// @brief Test whether p lies on the closed segment [a, b].
template <class T>
bool on_segment(const point<T> &p, const point<T> &a, const point<T> &b,
const T &epsilon = T{}) {
if (orientation(a, b, p, epsilon) != 0) {
return false;
}
return std::min(a.x, b.x) - epsilon <= p.x &&
p.x <= std::max(a.x, b.x) + epsilon &&
std::min(a.y, b.y) - epsilon <= p.y &&
p.y <= std::max(a.y, b.y) + epsilon;
}
/// @brief Test whether the closed segments [a, b] and [c, d] intersect.
template <class T>
bool segments_intersect(const point<T> &a, const point<T> &b, const point<T> &c,
const point<T> &d, const T &epsilon = T{}) {
int ab_c = orientation(a, b, c, epsilon);
int ab_d = orientation(a, b, d, epsilon);
int cd_a = orientation(c, d, a, epsilon);
int cd_b = orientation(c, d, b, epsilon);
if (ab_c == 0 && on_segment(c, a, b, epsilon)) {
return true;
}
if (ab_d == 0 && on_segment(d, a, b, epsilon)) {
return true;
}
if (cd_a == 0 && on_segment(a, c, d, epsilon)) {
return true;
}
if (cd_b == 0 && on_segment(b, c, d, epsilon)) {
return true;
}
return ab_c * ab_d < 0 && cd_a * cd_b < 0;
}
/// @brief Intersect the infinite lines through (a, b) and (c, d), returning
/// nullopt when they are parallel or coincident.
template <class T>
std::optional<point<long double>>
line_intersection(const point<T> &a, const point<T> &b, const point<T> &c,
const point<T> &d, long double epsilon = 0) {
point<long double> first{static_cast<long double>(a.x),
static_cast<long double>(a.y)};
point<long double> second{static_cast<long double>(b.x),
static_cast<long double>(b.y)};
point<long double> third{static_cast<long double>(c.x),
static_cast<long double>(c.y)};
point<long double> fourth{static_cast<long double>(d.x),
static_cast<long double>(d.y)};
point<long double> direction_a = second - first;
point<long double> direction_b = fourth - third;
long double denominator = cross(direction_a, direction_b);
if (std::abs(denominator) <= epsilon) {
return std::nullopt;
}
long double ratio = cross(third - first, direction_b) / denominator;
return first + direction_a * ratio;
}
/// @brief Return the convex hull in counter-clockwise order without repetition.
template <class T>
std::vector<point<T>> convex_hull(std::vector<point<T>> points,
bool keep_collinear = false) {
std::sort(points.begin(), points.end());
points.erase(std::unique(points.begin(), points.end()), points.end());
if (points.size() <= 1) {
return points;
}
bool all_collinear = true;
for (int i = 2; i < int(points.size()); i++) {
all_collinear &= orientation(points[0], points[1], points[i]) == 0;
}
if (keep_collinear && all_collinear) {
return points;
}
std::vector<point<T>> lower, upper;
for (const point<T> &p : points) {
while (lower.size() >= 2) {
int turn = orientation(lower[lower.size() - 2], lower.back(), p);
if (turn > 0 || (keep_collinear && turn == 0)) {
break;
}
lower.pop_back();
}
lower.push_back(p);
}
for (auto it = points.rbegin(); it != points.rend(); ++it) {
while (upper.size() >= 2) {
int turn = orientation(upper[upper.size() - 2], upper.back(), *it);
if (turn > 0 || (keep_collinear && turn == 0)) {
break;
}
upper.pop_back();
}
upper.push_back(*it);
}
lower.pop_back();
upper.pop_back();
lower.insert(lower.end(), upper.begin(), upper.end());
return lower;
}
/// @brief Return twice the signed area of a polygon.
template <class T> T polygon_area2(const std::vector<point<T>> &polygon) {
T result{};
for (int i = 0; i < int(polygon.size()); i++) {
result += cross(polygon[i], polygon[(i + 1) % polygon.size()]);
}
return result;
}
/// @brief Classify a point relative to a polygon: -1 outside, 0 boundary, 1
/// inside.
template <class T>
int point_in_polygon(const point<T> &p, const std::vector<point<T>> &polygon) {
bool inside = false;
for (int i = 0; i < int(polygon.size()); i++) {
point<T> a = polygon[i];
point<T> b = polygon[(i + 1) % polygon.size()];
if (on_segment(p, a, b)) {
return 0;
}
if (a.y <= p.y && p.y < b.y && orientation(a, b, p) > 0) {
inside = !inside;
}
if (b.y <= p.y && p.y < a.y && orientation(a, b, p) < 0) {
inside = !inside;
}
}
return inside ? 1 : -1;
}
/// @brief Return the Euclidean distance between two points.
template <class T> long double distance(const point<T> &a, const point<T> &b) {
return std::hypot(
static_cast<long double>(a.x) - static_cast<long double>(b.x),
static_cast<long double>(a.y) - static_cast<long double>(b.y));
}
} // namespace noya
namespace noya {
/// @brief Static two-dimensional KD-tree for rectangle reporting and nearest
/// neighbor queries.
template <class T, class Distance = std::conditional_t<std::is_integral_v<T>,
__int128_t, long double>>
struct kd_tree {
using point_type = point<T>;
using distance_type = Distance;
struct node {
int point_index = -1;
int left = -1;
int right = -1;
T min_x{};
T max_x{};
T min_y{};
T max_y{};
};
std::vector<point_type> points;
std::vector<node> nodes;
int root = -1;
kd_tree() = default;
explicit kd_tree(const std::vector<point_type> &input) { build(input); }
/// @brief Rebuild the balanced tree in O(n log n) expected time.
void build(const std::vector<point_type> &input) {
points = input;
nodes.clear();
root = -1;
std::vector<int> order(points.size());
std::iota(order.begin(), order.end(), 0);
if (!order.empty()) {
nodes.reserve(order.size());
root = build_range(order, 0, int(order.size()), 0);
}
}
/// @brief Return the number of stored points.
int size() const { return int(points.size()); }
/// @brief Return point indices in [left, right) x [bottom, top), stopping
/// after max_count results when max_count is nonnegative.
std::vector<int> range_query(const T &left, const T &right, const T &bottom,
const T &top, int max_count = -1) const {
assert(!(right < left) && !(top < bottom));
int limit = max_count < 0 ? size() : max_count;
std::vector<int> result;
result.reserve(std::min(size(), limit));
range_query_at(root, left, right, bottom, top, limit, result);
return result;
}
/// @brief Return (point index, squared distance) for a nearest point; the
/// index is -1 when the tree is empty.
std::pair<int, distance_type> nearest(const point_type &query) const {
int best_index = -1;
distance_type best_distance{};
nearest_at(root, query, best_index, best_distance);
return {best_index, best_distance};
}
private:
static distance_type square_difference(const T &first, const T &second) {
distance_type difference =
static_cast<distance_type>(first) - static_cast<distance_type>(second);
return difference * difference;
}
distance_type point_distance(int index, const point_type &query) const {
return square_difference(points[index].x, query.x) +
square_difference(points[index].y, query.y);
}
distance_type box_distance(const node ¤t,
const point_type &query) const {
distance_type result{};
if (query.x < current.min_x) {
result += square_difference(query.x, current.min_x);
} else if (current.max_x < query.x) {
result += square_difference(query.x, current.max_x);
}
if (query.y < current.min_y) {
result += square_difference(query.y, current.min_y);
} else if (current.max_y < query.y) {
result += square_difference(query.y, current.max_y);
}
return result;
}
int build_range(std::vector<int> &order, int left, int right, int depth) {
int middle = (left + right) / 2;
bool split_x = (depth & 1) == 0;
auto compare = [&](int first, int second) {
if (split_x) {
if (points[first].x != points[second].x) {
return points[first].x < points[second].x;
}
if (points[first].y != points[second].y) {
return points[first].y < points[second].y;
}
} else {
if (points[first].y != points[second].y) {
return points[first].y < points[second].y;
}
if (points[first].x != points[second].x) {
return points[first].x < points[second].x;
}
}
return first < second;
};
std::nth_element(order.begin() + left, order.begin() + middle,
order.begin() + right, compare);
int id = int(nodes.size());
int point_index = order[middle];
nodes.push_back({point_index, -1, -1, points[point_index].x,
points[point_index].x, points[point_index].y,
points[point_index].y});
if (left < middle) {
nodes[id].left = build_range(order, left, middle, depth + 1);
extend_box(nodes[id], nodes[nodes[id].left]);
}
if (middle + 1 < right) {
nodes[id].right = build_range(order, middle + 1, right, depth + 1);
extend_box(nodes[id], nodes[nodes[id].right]);
}
return id;
}
static void extend_box(node &target, const node &source) {
target.min_x = std::min(target.min_x, source.min_x);
target.max_x = std::max(target.max_x, source.max_x);
target.min_y = std::min(target.min_y, source.min_y);
target.max_y = std::max(target.max_y, source.max_y);
}
static bool box_disjoint(const node ¤t, const T &left, const T &right,
const T &bottom, const T &top) {
return current.max_x < left || !(current.min_x < right) ||
current.max_y < bottom || !(current.min_y < top);
}
static bool box_contained(const node ¤t, const T &left, const T &right,
const T &bottom, const T &top) {
return !(current.min_x < left) && current.max_x < right &&
!(current.min_y < bottom) && current.max_y < top;
}
void collect_subtree(int id, int limit, std::vector<int> &result) const {
if (id == -1 || int(result.size()) == limit) {
return;
}
result.push_back(nodes[id].point_index);
collect_subtree(nodes[id].left, limit, result);
collect_subtree(nodes[id].right, limit, result);
}
void range_query_at(int id, const T &left, const T &right, const T &bottom,
const T &top, int limit, std::vector<int> &result) const {
if (id == -1 || int(result.size()) == limit ||
box_disjoint(nodes[id], left, right, bottom, top)) {
return;
}
if (box_contained(nodes[id], left, right, bottom, top)) {
collect_subtree(id, limit, result);
return;
}
int index = nodes[id].point_index;
const point_type &value = points[index];
if (!(value.x < left) && value.x < right && !(value.y < bottom) &&
value.y < top) {
result.push_back(index);
}
range_query_at(nodes[id].left, left, right, bottom, top, limit, result);
range_query_at(nodes[id].right, left, right, bottom, top, limit, result);
}
void nearest_at(int id, const point_type &query, int &best_index,
distance_type &best_distance) const {
if (id == -1) {
return;
}
int index = nodes[id].point_index;
distance_type distance = point_distance(index, query);
if (best_index == -1 || distance < best_distance ||
(distance == best_distance && index < best_index)) {
best_index = index;
best_distance = distance;
}
int first = nodes[id].left;
int second = nodes[id].right;
distance_type first_bound =
first == -1 ? distance_type{} : box_distance(nodes[first], query);
distance_type second_bound =
second == -1 ? distance_type{} : box_distance(nodes[second], query);
if (first == -1 || (second != -1 && second_bound < first_bound)) {
std::swap(first, second);
std::swap(first_bound, second_bound);
}
if (first != -1 && first_bound <= best_distance) {
nearest_at(first, query, best_index, best_distance);
}
if (second != -1 && second_bound <= best_distance) {
nearest_at(second, query, best_index, best_distance);
}
}
};
} // namespace noya