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

OpenTTD Source 14.0-beta3

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