KdTree-Impl.h
Go to the documentation of this file.
1 /*************************************************************************
2 > File Name: KdTree-Impl.h
3 > Project Name: CubbyFlow
4 > Author: Chan-Ho Chris Ohk
5 > Purpose: Generic k-d tree structure.
6 > Created Time: 2017/12/05
7 > Copyright (c) 2018, Chan-Ho Chris Ohk
8 *************************************************************************/
9 #ifndef CUBBYFLOW_KDTREE_IMPL_H
10 #define CUBBYFLOW_KDTREE_IMPL_H
11 
12 #include <numeric>
13 
14 namespace CubbyFlow
15 {
16  template <typename T, size_t K>
18  {
19  child = std::numeric_limits<size_t>::max();
20  }
21 
22  template <typename T, size_t K>
23  void KdTree<T, K>::Node::InitLeaf(size_t it, const Point& pt)
24  {
25  flags = K;
26  item = it;
27  child = std::numeric_limits<size_t>::max();
28  point = pt;
29  }
30 
31  template <typename T, size_t K>
32  void KdTree<T, K>::Node::InitInternal(size_t axis, size_t it, size_t c, const Point& pt)
33  {
34  flags = axis;
35  item = it;
36  child = c;
37  point = pt;
38  }
39 
40  template <typename T, size_t K>
42  {
43  return flags == K;
44  }
45 
46  template <typename T, size_t K>
48  {
49  // Do nothing
50  }
51 
52  template <typename T, size_t K>
54  {
55  m_points.resize(points.size());
56  std::copy(points.begin(), points.end(), m_points.begin());
57 
58  if (m_points.empty())
59  {
60  return;
61  }
62 
63  m_nodes.clear();
64 
65  std::vector<size_t> itemIndices(m_points.size());
66  std::iota(std::begin(itemIndices), std::end(itemIndices), 0);
67 
68  Build(0, itemIndices.data(), m_points.size(), 0);
69  }
70 
71  template <typename T, size_t K>
73  const Point& origin, T radius,
74  const std::function<void(size_t, const Point&)>& callback) const
75  {
76  const T r2 = radius * radius;
77 
78  // prepare to traverse the tree for sphere
79  static const int maxTreeDepth = 8 * sizeof(size_t);
80  const Node* todo[maxTreeDepth];
81  size_t todoPos = 0;
82 
83  // traverse the tree nodes for sphere
84  const Node* node = m_nodes.data();
85 
86  while (node != nullptr)
87  {
88  if (node->item != std::numeric_limits<size_t>::max() && (node->point - origin).LengthSquared() <= r2)
89  {
90  callback(node->item, node->point);
91  }
92 
93  if (node->IsLeaf())
94  {
95  // grab next node to process from todo stack
96  if (todoPos > 0)
97  {
98  // dequeue
99  --todoPos;
100  node = todo[todoPos];
101  }
102  else
103  {
104  break;
105  }
106  }
107  else
108  {
109  // get node children pointers for sphere
110  const Node* firstChild = node + 1;
111  const Node* secondChild = const_cast<Node*>(&m_nodes[node->child]);
112 
113  // advance to next child node, possibly enqueue other child
114  const size_t axis = node->flags;
115  const T plane = node->point[axis];
116 
117  if (plane - origin[axis] > radius)
118  {
119  node = firstChild;
120  }
121  else if (origin[axis] - plane > radius)
122  {
123  node = secondChild;
124  }
125  else
126  {
127  // enqueue secondChild in todo stack
128  todo[todoPos] = secondChild;
129  ++todoPos;
130  node = firstChild;
131  }
132  }
133  }
134  }
135 
136  template <typename T, size_t K>
137  bool KdTree<T, K>::HasNearbyPoint(const Point& origin, T radius) const
138  {
139  const T r2 = radius * radius;
140 
141  // prepare to traverse the tree for sphere
142  static const int maxTreeDepth = 8 * sizeof(size_t);
143  const Node* todo[maxTreeDepth];
144  size_t todoPos = 0;
145 
146  // traverse the tree nodes for sphere
147  const Node* node = m_nodes.data();
148 
149  while (node != nullptr)
150  {
151  if (node->item != std::numeric_limits<size_t>::max() && (node->point - origin).LengthSquared() <= r2)
152  {
153  return true;
154  }
155 
156  if (node->IsLeaf())
157  {
158  // grab next node to process from todo stack
159  if (todoPos > 0)
160  {
161  // dequeue
162  --todoPos;
163  node = todo[todoPos];
164  }
165  else
166  {
167  break;
168  }
169  }
170  else
171  {
172  // get node children pointers for sphere
173  const Node* firstChild = node + 1;
174  const Node* secondChild = const_cast<Node*>(&m_nodes[node->child]);
175 
176  // advance to next child node, possibly enqueue other child
177  const size_t axis = node->flags;
178  const T plane = node->point[axis];
179 
180  if (origin[axis] < plane && plane - origin[axis] > radius)
181  {
182  node = firstChild;
183  }
184  else if (origin[axis] > plane && origin[axis] - plane > radius)
185  {
186  node = secondChild;
187  }
188  else
189  {
190  // enqueue secondChild in todo stack
191  todo[todoPos] = secondChild;
192  ++todoPos;
193  node = firstChild;
194  }
195  }
196  }
197 
198  return false;
199  }
200 
201  template <typename T, size_t K>
202  size_t KdTree<T, K>::GetNearestPoint(const Point& origin) const
203  {
204  // prepare to traverse the tree for sphere
205  static const int maxTreeDepth = 8 * sizeof(size_t);
206  const Node* todo[maxTreeDepth];
207  size_t todoPos = 0;
208 
209  // traverse the tree nodes for sphere
210  const Node* node = m_nodes.data();
211  size_t nearest = 0;
212  T minDist2 = (node->point - origin).LengthSquared();
213 
214  while (node != nullptr)
215  {
216  const T newDist2 = (node->point - origin).LengthSquared();
217  if (newDist2 <= minDist2)
218  {
219  nearest = node->item;
220  minDist2 = newDist2;
221  }
222 
223  if (node->IsLeaf())
224  {
225  // grab next node to process from todo stack
226  if (todoPos > 0)
227  {
228  // Dequeue
229  --todoPos;
230  node = todo[todoPos];
231  }
232  else
233  {
234  break;
235  }
236  }
237  else
238  {
239  // get node children pointers for sphere
240  const Node* firstChild = node + 1;
241  const Node* secondChild = static_cast<Node*>(&m_nodes[node->child]);
242 
243  // advance to next child node, possibly enqueue other child
244  const size_t axis = node->flags;
245  const T plane = node->point[axis];
246  const T minDist = std::sqrt(minDist2);
247 
248  if (plane - origin[axis] > minDist)
249  {
250  node = firstChild;
251  }
252  else if (origin[axis] - plane > minDist)
253  {
254  node = secondChild;
255  }
256  else
257  {
258  // enqueue secondChild in todo stack
259  todo[todoPos] = secondChild;
260  ++todoPos;
261  node = firstChild;
262  }
263  }
264  }
265 
266  return nearest;
267  }
268 
269  template <typename T, size_t K>
270  void KdTree<T, K>::Reserve(size_t numPoints, size_t numNodes)
271  {
272  m_points.resize(numPoints);
273  m_nodes.resize(numNodes);
274  }
275 
276  template <typename T, size_t K>
278  {
279  return m_points.begin();
280  };
281 
282  template <typename T, size_t K>
284  {
285  return m_points.end();
286  };
287 
288  template <typename T, size_t K>
290  {
291  return m_points.begin();
292  };
293 
294  template <typename T, size_t K>
296  {
297  return m_points.end();
298  };
299 
300  template <typename T, size_t K>
302  {
303  return m_nodes.begin();
304  };
305 
306  template <typename T, size_t K>
308  {
309  return m_nodes.end();
310  };
311 
312  template <typename T, size_t K>
314  {
315  return m_nodes.begin();
316  };
317 
318  template <typename T, size_t K>
320  {
321  return m_nodes.end();
322  };
323 
324  template <typename T, size_t K>
325  size_t KdTree<T, K>::Build(size_t nodeIndex, size_t* itemIndices, size_t nItems, size_t currentDepth)
326  {
327  // add a node
328  m_nodes.emplace_back();
329 
330  // initialize leaf node if termination criteria met
331  if (nItems == 0)
332  {
333  m_nodes[nodeIndex].InitLeaf(std::numeric_limits<size_t>::max(), {});
334  return currentDepth + 1;
335  }
336  if (nItems == 1)
337  {
338  m_nodes[nodeIndex].InitLeaf(itemIndices[0], m_points[itemIndices[0]]);
339  return currentDepth + 1;
340  }
341 
342  // choose which axis to split along
343  BBox nodeBound;
344  for (size_t i = 0; i < nItems; ++i)
345  {
346  nodeBound.Merge(m_points[itemIndices[i]]);
347  }
348  Point d = nodeBound.upperCorner - nodeBound.lowerCorner;
349  size_t axis = static_cast<size_t>(d.DominantAxis());
350 
351  // pick mid point
352  std::nth_element(itemIndices, itemIndices + nItems / 2, itemIndices + nItems,
353  [&](size_t a, size_t b)
354  {
355  return m_points[a][axis] < m_points[b][axis];
356  });
357  const size_t midPoint = nItems / 2;
358 
359  // recursively initialize children nodes
360  const size_t d0 = Build(nodeIndex + 1, itemIndices, midPoint, currentDepth + 1);
361  m_nodes[nodeIndex].InitInternal(axis, itemIndices[midPoint], m_nodes.size(), m_points[itemIndices[midPoint]]);
362  const size_t d1 = Build(m_nodes[nodeIndex].child, itemIndices + midPoint + 1, nItems - midPoint - 1, currentDepth + 1);
363 
364  return std::max(d0, d1);
365  }
366 }
367 
368 #endif
void Reserve(size_t numPoints, size_t numNodes)
Reserves memory space for this tree.
Definition: KdTree-Impl.h:270
NodeIterator EndNode()
Returns the mutable end iterator of the node.
Definition: KdTree-Impl.h:307
typename ContainerType::const_iterator ConstIterator
Definition: KdTree.h:56
typename NodeContainerType::iterator NodeIterator
Definition: KdTree.h:59
const T *const begin() const
Returns the begin iterator of the array.
Definition: ArrayAccessor1-Impl.h:207
Iterator begin()
Returns the mutable begin iterator of the item.
Definition: KdTree-Impl.h:277
bool IsLeaf() const
Returns true if leaf.
Definition: KdTree-Impl.h:41
Node()
Default contructor.
Definition: KdTree-Impl.h:17
bool HasNearbyPoint(const Point &origin, T radius) const
Definition: KdTree-Impl.h:137
1-D read-only array accessor class.
Definition: ArrayAccessor1.h:185
NodeIterator BeginNode()
Returns the mutable begin iterator of the node.
Definition: KdTree-Impl.h:301
void InitInternal(size_t axis, size_t it, size_t c, const Point &pt)
Initializes internal node.
Definition: KdTree-Impl.h:32
size_t GetNearestPoint(const Point &origin) const
Returns index of the nearest point.
Definition: KdTree-Impl.h:202
BoundingBox< T, K > BBox
Definition: KdTree.h:23
Definition: pybind11Utils.h:24
size_t size() const
Returns size of the array.
Definition: ArrayAccessor1-Impl.h:219
typename NodeContainerType::const_iterator ConstNodeIterator
Definition: KdTree.h:60
void Build(const ConstArrayAccessor1< Point > &points)
Builds internal acceleration structure for given points list.
Definition: KdTree-Impl.h:53
const T *const end() const
Returns the end iterator of the array.
Definition: ArrayAccessor1-Impl.h:213
size_t child
Right child index. Note that left child index is this node index + 1.
Definition: KdTree.h:33
Iterator end()
Returns the mutable end iterator of the item.
Definition: KdTree-Impl.h:283
KdTree()
Constructs an empty kD-tree instance.
Definition: KdTree-Impl.h:47
typename ContainerType::iterator Iterator
Definition: KdTree.h:55
void ForEachNearbyPoint(const Point &origin, T radius, const std::function< void(size_t, const Point &)> &callback) const
Definition: KdTree-Impl.h:72
void InitLeaf(size_t it, const Point &pt)
Initializes leaf node.
Definition: KdTree-Impl.h:23
Vector< T, K > Point
Definition: KdTree.h:22