OpenTTD Source: src/core/kdtree.hpp Source File

OpenTTD Source 15.0-RC2

kdtree.hpp
Go to the documentation of this file.
1/*
2 * This file is part of OpenTTD.
3 * OpenTTD is free software; you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, version 2.
4 * OpenTTD is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
5 * See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with OpenTTD. If not, see <https://www.gnu.org/licenses/old-licenses/gpl-2.0>.
6 */
7
10#ifndef KDTREE_HPP
11#define KDTREE_HPP
12
32template <typename T, typename TxyFunc, typename CoordT, typename DistT>
33 class Kdtree {
35 struct node {
36 T element;
37 size_t left;
38 size_t right;
39
40 node(T element) : element(element), left(INVALID_NODE), right(INVALID_NODE) { }
41 };
42
43 static const size_t INVALID_NODE = SIZE_MAX;
44 static const size_t MIN_REBALANCE_THRESHOLD = 8;
45
46 std::vector<node> nodes;
47 std::vector<size_t> free_list;
48 size_t root;
49 size_t unbalanced;
50
52 size_t AddNode(const T &element)
53 {
54 if (this->free_list.empty()) {
55 this->nodes.emplace_back(element);
56 return this->nodes.size() - 1;
57 } else {
58 size_t newidx = this->free_list.back();
59 this->free_list.pop_back();
60 this->nodes[newidx] = node{ element };
61 return newidx;
62 }
63 }
64
66 template <typename It>
67 CoordT SelectSplitCoord(It begin, It end, int level)
68 {
69 It mid = begin + (end - begin) / 2;
70 std::nth_element(begin, mid, end, [&](T a, T b) { return TxyFunc()(a, level % 2) < TxyFunc()(b, level % 2); });
71 return TxyFunc()(*mid, level % 2);
72 }
73
75 template <typename It>
76 size_t BuildSubtree(It begin, It end, int level)
77 {
78 ptrdiff_t count = end - begin;
79
80 if (count == 0) {
81 return INVALID_NODE;
82 } else if (count == 1) {
83 return this->AddNode(*begin);
84 } else if (count > 1) {
85 CoordT split_coord = this->SelectSplitCoord(begin, end, level);
86 It split = std::partition(begin, end, [&](T v) { return TxyFunc()(v, level % 2) < split_coord; });
87 size_t newidx = this->AddNode(*split);
88 this->nodes[newidx].left = this->BuildSubtree(begin, split, level + 1);
89 this->nodes[newidx].right = this->BuildSubtree(split + 1, end, level + 1);
90 return newidx;
91 } else {
92 NOT_REACHED();
93 }
94 }
95
97 bool Rebuild(const T *include_element, const T *exclude_element)
98 {
99 size_t initial_count = this->Count();
100 if (initial_count < MIN_REBALANCE_THRESHOLD) return false;
101
102 T root_element = this->nodes[this->root].element;
103 std::vector<T> elements = this->FreeSubtree(this->root);
104 elements.push_back(root_element);
105
106 if (include_element != nullptr) {
107 elements.push_back(*include_element);
108 initial_count++;
109 }
110 if (exclude_element != nullptr) {
111 typename std::vector<T>::iterator removed = std::remove(elements.begin(), elements.end(), *exclude_element);
112 elements.erase(removed, elements.end());
113 initial_count--;
114 }
115
116 this->Build(elements.begin(), elements.end());
117 assert(initial_count == this->Count());
118 return true;
119 }
120
122 void InsertRecursive(const T &element, size_t node_idx, int level)
123 {
124 /* Dimension index of current level */
125 int dim = level % 2;
126 /* Node reference */
127 node &n = this->nodes[node_idx];
128
129 /* Coordinate of element splitting at this node */
130 CoordT nc = TxyFunc()(n.element, dim);
131 /* Coordinate of the new element */
132 CoordT ec = TxyFunc()(element, dim);
133 /* Which side to insert on */
134 size_t &next = (ec < nc) ? n.left : n.right;
135
136 if (next == INVALID_NODE) {
137 /* New leaf */
138 size_t newidx = this->AddNode(element);
139 /* Vector may have been reallocated at this point, n and next are invalid */
140 node &nn = this->nodes[node_idx];
141 if (ec < nc) nn.left = newidx; else nn.right = newidx;
142 } else {
143 this->InsertRecursive(element, next, level + 1);
144 }
145 }
146
151 std::vector<T> FreeSubtree(size_t node_idx)
152 {
153 std::vector<T> subtree_elements;
154 node &n = this->nodes[node_idx];
155
156 /* We'll be appending items to the free_list, get index of our first item */
157 size_t first_free = this->free_list.size();
158 /* Prepare the descent with our children */
159 if (n.left != INVALID_NODE) this->free_list.push_back(n.left);
160 if (n.right != INVALID_NODE) this->free_list.push_back(n.right);
161 n.left = n.right = INVALID_NODE;
162
163 /* Recursively free the nodes being collected */
164 for (size_t i = first_free; i < this->free_list.size(); i++) {
165 node &fn = this->nodes[this->free_list[i]];
166 subtree_elements.push_back(fn.element);
167 if (fn.left != INVALID_NODE) this->free_list.push_back(fn.left);
168 if (fn.right != INVALID_NODE) this->free_list.push_back(fn.right);
169 fn.left = fn.right = INVALID_NODE;
170 }
171
172 return subtree_elements;
173 }
174
182 size_t RemoveRecursive(const T &element, size_t node_idx, int level)
183 {
184 /* Node reference */
185 node &n = this->nodes[node_idx];
186
187 if (n.element == element) {
188 /* Remove this one */
189 this->free_list.push_back(node_idx);
190 if (n.left == INVALID_NODE && n.right == INVALID_NODE) {
191 /* Simple case, leaf, new child node for parent is "none" */
192 return INVALID_NODE;
193 } else {
194 /* Complex case, rebuild the sub-tree */
195 std::vector<T> subtree_elements = this->FreeSubtree(node_idx);
196 return this->BuildSubtree(subtree_elements.begin(), subtree_elements.end(), level);;
197 }
198 } else {
199 /* Search in a sub-tree */
200 /* Dimension index of current level */
201 int dim = level % 2;
202 /* Coordinate of element splitting at this node */
203 CoordT nc = TxyFunc()(n.element, dim);
204 /* Coordinate of the element being removed */
205 CoordT ec = TxyFunc()(element, dim);
206 /* Which side to remove from */
207 size_t next = (ec < nc) ? n.left : n.right;
208 assert(next != INVALID_NODE); // node must exist somewhere and must be found before a leaf is reached
209 /* Descend */
210 size_t new_branch = this->RemoveRecursive(element, next, level + 1);
211 if (new_branch != next) {
212 /* Vector may have been reallocated at this point, n and next are invalid */
213 node &nn = this->nodes[node_idx];
214 if (ec < nc) nn.left = new_branch; else nn.right = new_branch;
215 }
216 return node_idx;
217 }
218 }
219
220
221 DistT ManhattanDistance(const T &element, CoordT x, CoordT y) const
222 {
223 return abs((DistT)TxyFunc()(element, 0) - (DistT)x) + abs((DistT)TxyFunc()(element, 1) - (DistT)y);
224 }
225
227 using node_distance = std::pair<T, DistT>;
229 static node_distance SelectNearestNodeDistance(const node_distance &a, const node_distance &b)
230 {
231 if (a.second < b.second) return a;
232 if (b.second < a.second) return b;
233 if (a.first < b.first) return a;
234 if (b.first < a.first) return b;
235 NOT_REACHED(); // a.first == b.first: same element must not be inserted twice
236 }
238 node_distance FindNearestRecursive(CoordT xy[2], size_t node_idx, int level, DistT limit = std::numeric_limits<DistT>::max()) const
239 {
240 /* Dimension index of current level */
241 int dim = level % 2;
242 /* Node reference */
243 const node &n = this->nodes[node_idx];
244
245 /* Coordinate of element splitting at this node */
246 CoordT c = TxyFunc()(n.element, dim);
247 /* This node's distance to target */
248 DistT thisdist = this->ManhattanDistance(n.element, xy[0], xy[1]);
249 /* Assume this node is the best choice for now */
250 node_distance best = std::make_pair(n.element, thisdist);
251
252 /* Next node to visit */
253 size_t next = (xy[dim] < c) ? n.left : n.right;
254 if (next != INVALID_NODE) {
255 /* Check if there is a better node down the tree */
256 best = SelectNearestNodeDistance(best, this->FindNearestRecursive(xy, next, level + 1));
257 }
258
259 limit = std::min(best.second, limit);
260
261 /* Check if the distance from current best is worse than distance from target to splitting line,
262 * if it is we also need to check the other side of the split. */
263 size_t opposite = (xy[dim] >= c) ? n.left : n.right; // reverse of above
264 if (opposite != INVALID_NODE && limit >= abs((int)xy[dim] - (int)c)) {
265 node_distance other_candidate = this->FindNearestRecursive(xy, opposite, level + 1, limit);
266 best = SelectNearestNodeDistance(best, other_candidate);
267 }
268
269 return best;
270 }
271
272 template <typename Outputter>
273 void FindContainedRecursive(CoordT p1[2], CoordT p2[2], size_t node_idx, int level, const Outputter &outputter) const
274 {
275 /* Dimension index of current level */
276 int dim = level % 2;
277 /* Node reference */
278 const node &n = this->nodes[node_idx];
279
280 /* Coordinate of element splitting at this node */
281 CoordT ec = TxyFunc()(n.element, dim);
282 /* Opposite coordinate of element */
283 CoordT oc = TxyFunc()(n.element, 1 - dim);
284
285 /* Test if this element is within rectangle */
286 if (ec >= p1[dim] && ec < p2[dim] && oc >= p1[1 - dim] && oc < p2[1 - dim]) outputter(n.element);
287
288 /* Recurse left if part of rectangle is left of split */
289 if (p1[dim] < ec && n.left != INVALID_NODE) this->FindContainedRecursive(p1, p2, n.left, level + 1, outputter);
290
291 /* Recurse right if part of rectangle is right of split */
292 if (p2[dim] > ec && n.right != INVALID_NODE) this->FindContainedRecursive(p1, p2, n.right, level + 1, outputter);
293 }
294
296 size_t CountValue(const T &element, size_t node_idx) const
297 {
298 if (node_idx == INVALID_NODE) return 0;
299 const node &n = this->nodes[node_idx];
300 return this->CountValue(element, n.left) + this->CountValue(element, n.right) + ((n.element == element) ? 1 : 0);
301 }
302
303 void IncrementUnbalanced(size_t amount = 1)
304 {
305 this->unbalanced += amount;
306 }
307
309 bool IsUnbalanced() const
310 {
311 size_t count = this->Count();
312 if (count < MIN_REBALANCE_THRESHOLD) return false;
313 return this->unbalanced > count / 4;
314 }
315
317 void CheckInvariant(size_t node_idx, int level, CoordT min_x, CoordT max_x, CoordT min_y, CoordT max_y) const
318 {
319 if (node_idx == INVALID_NODE) return;
320
321 const node &n = this->nodes[node_idx];
322 CoordT cx = TxyFunc()(n.element, 0);
323 CoordT cy = TxyFunc()(n.element, 1);
324
325 assert(cx >= min_x);
326 assert(cx < max_x);
327 assert(cy >= min_y);
328 assert(cy < max_y);
329
330 if (level % 2 == 0) {
331 /* split in dimension 0 = x */
332 this->CheckInvariant(n.left, level + 1, min_x, cx, min_y, max_y);
333 this->CheckInvariant(n.right, level + 1, cx, max_x, min_y, max_y);
334 } else {
335 /* split in dimension 1 = y */
336 this->CheckInvariant(n.left, level + 1, min_x, max_x, min_y, cy);
337 this->CheckInvariant(n.right, level + 1, min_x, max_x, cy, max_y);
338 }
339 }
340
342 void CheckInvariant() const
343 {
344#ifdef KDTREE_DEBUG
345 this->CheckInvariant(this->root, 0, std::numeric_limits<CoordT>::min(), std::numeric_limits<CoordT>::max(), std::numeric_limits<CoordT>::min(), std::numeric_limits<CoordT>::max());
346#endif
347 }
348
349public:
351 Kdtree() : root(INVALID_NODE), unbalanced(0) { }
352
359 template <typename It>
360 void Build(It begin, It end)
361 {
362 this->nodes.clear();
363 this->free_list.clear();
364 this->unbalanced = 0;
365 if (begin == end) return;
366 this->nodes.reserve(end - begin);
367
368 this->root = this->BuildSubtree(begin, end, 0);
369 this->CheckInvariant();
370 }
371
375 void Clear()
376 {
377 this->nodes.clear();
378 this->free_list.clear();
379 this->unbalanced = 0;
380 return;
381 }
382
386 void Rebuild()
387 {
388 this->Rebuild(nullptr, nullptr);
389 }
390
396 void Insert(const T &element)
397 {
398 if (this->Count() == 0) {
399 this->root = this->AddNode(element);
400 } else {
401 if (!this->IsUnbalanced() || !this->Rebuild(&element, nullptr)) {
402 this->InsertRecursive(element, this->root, 0);
403 this->IncrementUnbalanced();
404 }
405 this->CheckInvariant();
406 }
407 }
408
415 void Remove(const T &element)
416 {
417 size_t count = this->Count();
418 if (count == 0) return;
419 if (!this->IsUnbalanced() || !this->Rebuild(nullptr, &element)) {
420 /* If the removed element is the root node, this modifies this->root */
421 this->root = this->RemoveRecursive(element, this->root, 0);
422 this->IncrementUnbalanced();
423 }
424 this->CheckInvariant();
425 }
426
428 size_t Count() const
429 {
430 assert(this->free_list.size() <= this->nodes.size());
431 return this->nodes.size() - this->free_list.size();
432 }
433
439 T FindNearest(CoordT x, CoordT y) const
440 {
441 assert(this->Count() > 0);
442
443 CoordT xy[2] = { x, y };
444 return this->FindNearestRecursive(xy, this->root, 0).first;
445 }
446
456 template <typename Outputter>
457 void FindContained(CoordT x1, CoordT y1, CoordT x2, CoordT y2, const Outputter &outputter) const
458 {
459 assert(x1 < x2);
460 assert(y1 < y2);
461
462 if (this->Count() == 0) return;
463
464 CoordT p1[2] = { x1, y1 };
465 CoordT p2[2] = { x2, y2 };
466 this->FindContainedRecursive(p1, p2, this->root, 0, outputter);
467 }
468
473 std::vector<T> FindContained(CoordT x1, CoordT y1, CoordT x2, CoordT y2) const
474 {
475 std::vector<T> result;
476 this->FindContained(x1, y1, x2, y2, [&result](T e) {result.push_back(e); });
477 return result;
478 }
479};
480
481#endif
Kdtree
K-dimensional tree, specialised for 2-dimensional space.
Definition kdtree.hpp:33
Kdtree::MIN_REBALANCE_THRESHOLD
static const size_t MIN_REBALANCE_THRESHOLD
Arbitrary value for "not worth rebalancing".
Definition kdtree.hpp:44
Kdtree::Build
void Build(It begin, It end)
Clear and rebuild the tree from a new sequence of elements,.
Definition kdtree.hpp:360
Kdtree::Count
size_t Count() const
Get number of elements stored in tree.
Definition kdtree.hpp:428
Kdtree::IsUnbalanced
bool IsUnbalanced() const
Check if the entire tree is in need of rebuilding.
Definition kdtree.hpp:309
Kdtree::FindContained
void FindContained(CoordT x1, CoordT y1, CoordT x2, CoordT y2, const Outputter &outputter) const
Find all items contained within the given rectangle.
Definition kdtree.hpp:457
Kdtree::Rebuild
bool Rebuild(const T *include_element, const T *exclude_element)
Rebuild the tree with all existing elements, optionally adding or removing one more.
Definition kdtree.hpp:97
Kdtree::free_list
std::vector< size_t > free_list
List of dead indices in the nodes vector.
Definition kdtree.hpp:47
Kdtree::FindNearestRecursive
node_distance FindNearestRecursive(CoordT xy[2], size_t node_idx, int level, DistT limit=std::numeric_limits< DistT >::max()) const
Search a sub-tree for the element nearest to a given point.
Definition kdtree.hpp:238
Kdtree::InsertRecursive
void InsertRecursive(const T &element, size_t node_idx, int level)
Insert one element in the tree somewhere below node_idx.
Definition kdtree.hpp:122
Kdtree::Rebuild
void Rebuild()
Reconstruct the tree with the same elements, letting it be fully balanced.
Definition kdtree.hpp:386
Kdtree::node_distance
std::pair< T, DistT > node_distance
A data element and its distance to a searched-for point.
Definition kdtree.hpp:227
Kdtree::Insert
void Insert(const T &element)
Insert a single element in the tree.
Definition kdtree.hpp:396
Kdtree::Kdtree
Kdtree()
Construct a new Kdtree with the given xyfunc.
Definition kdtree.hpp:351
Kdtree::nodes
std::vector< node > nodes
Pool of all nodes in the tree.
Definition kdtree.hpp:46
Kdtree::FindContained
std::vector< T > FindContained(CoordT x1, CoordT y1, CoordT x2, CoordT y2) const
Find all items contained within the given rectangle.
Definition kdtree.hpp:473
Kdtree::SelectNearestNodeDistance
static node_distance SelectNearestNodeDistance(const node_distance &a, const node_distance &b)
Ordering function for node_distance objects, elements with equal distance are ordered by less-than co...
Definition kdtree.hpp:229
Kdtree::CheckInvariant
void CheckInvariant(size_t node_idx, int level, CoordT min_x, CoordT max_x, CoordT min_y, CoordT max_y) const
Verify that the invariant is true for a sub-tree, assert if not.
Definition kdtree.hpp:317
Kdtree::Remove
void Remove(const T &element)
Remove a single element from the tree, if it exists.
Definition kdtree.hpp:415
Kdtree::INVALID_NODE
static const size_t INVALID_NODE
Index value indicating no-such-node.
Definition kdtree.hpp:43
Kdtree::root
size_t root
Index of root node.
Definition kdtree.hpp:48
Kdtree::RemoveRecursive
size_t RemoveRecursive(const T &element, size_t node_idx, int level)
Find and remove one element from the tree.
Definition kdtree.hpp:182
Kdtree::BuildSubtree
size_t BuildSubtree(It begin, It end, int level)
Construct a subtree from elements between begin and end iterators, return index of root.
Definition kdtree.hpp:76
Kdtree::AddNode
size_t AddNode(const T &element)
Create one new node in the tree, return its index in the pool.
Definition kdtree.hpp:52
Kdtree::CheckInvariant
void CheckInvariant() const
Verify the invariant for the entire tree, does nothing unless KDTREE_DEBUG is defined.
Definition kdtree.hpp:342
Kdtree::FindNearest
T FindNearest(CoordT x, CoordT y) const
Find the element closest to given coordinate, in Manhattan distance.
Definition kdtree.hpp:439
Kdtree::SelectSplitCoord
CoordT SelectSplitCoord(It begin, It end, int level)
Find a coordinate value to split a range of elements at.
Definition kdtree.hpp:67
Kdtree::FreeSubtree
std::vector< T > FreeSubtree(size_t node_idx)
Free all children of the given node.
Definition kdtree.hpp:151
Kdtree::Clear
void Clear()
Clear the tree.
Definition kdtree.hpp:375
Kdtree::CountValue
size_t CountValue(const T &element, size_t node_idx) const
Debugging function, counts number of occurrences of an element regardless of its correct position in ...
Definition kdtree.hpp:296
Kdtree::unbalanced
size_t unbalanced
Number approximating how unbalanced the tree might be.
Definition kdtree.hpp:49
abs
constexpr T abs(const T a)
Returns the absolute value of (scalar) variable.
Definition math_func.hpp:23
Kdtree::node
Type of a node in the tree.
Definition kdtree.hpp:35
Kdtree::node::element
T element
Element stored at node.
Definition kdtree.hpp:36
Kdtree::node::left
size_t left
Index of node to the left, INVALID_NODE if none.
Definition kdtree.hpp:37
Kdtree::node::right
size_t right
Index of node to the right, INVALID_NODE if none.
Definition kdtree.hpp:38