void recursiveSearch(uint32_t nodeIndex, const Vec3f& point, float& maxDistanceSquared, ProcessFunctor process) const {
const KdNode& node = m_Nodes[nodeIndex];
uint32_t axis = node.m_nSplitAxis;
if(axis != 3) {
float distSquared = sqr(point[axis] - node.m_fSplitPosition);
if(point[axis] <= node.m_fSplitPosition) {
if(node.m_bHasLeftChild) {
recursiveSearch(nodeIndex + 1, point, maxDistanceSquared, process);
}
if(distSquared < maxDistanceSquared && node.m_nRightChildIndex < m_Nodes.size()) {
recursiveSearch(node.m_nRightChildIndex, point, maxDistanceSquared, process);
}
} else {
if(node.m_nRightChildIndex < m_Nodes.size()) {
recursiveSearch(node.m_nRightChildIndex, point, maxDistanceSquared, process);
}
if(distSquared < maxDistanceSquared && node.m_bHasLeftChild) {
recursiveSearch(nodeIndex + 1, point, maxDistanceSquared, process);
}
}
}
float distSquared = sqr_distance(m_NodesData[nodeIndex].m_Position, point);
if(distSquared < maxDistanceSquared) {
process(m_NodesData[nodeIndex].m_nIndex, m_NodesData[nodeIndex].m_Position, distSquared, maxDistanceSquared);
}
}
uint32_t recursiveSearchNearestNeighbour(uint32_t nodeIndex, const Vec3f& point,
float& distSquared, Predicate predicate) const {
auto node = m_Nodes[nodeIndex];
uint32_t index = m_NodesData[nodeIndex].m_nIndex;
float candidateDistSquared = sqr_distance(point, m_NodesData[nodeIndex].m_Position);
// If the node is a leaf, end of the recursion
if(node.m_nSplitAxis == 3) {
if(candidateDistSquared < distSquared && predicate(index)) {
distSquared = candidateDistSquared;
return nodeIndex;
}
return KdNode::NO_NODE;
}
// If the node is inner, find a neighbour in the side of the point
uint32_t currentBestMatch = KdNode::NO_NODE;
uint8_t axis = node.m_nSplitAxis;
bool isAtLeft = (point[axis] <= node.m_fSplitPosition);
// Left side
if(isAtLeft && node.m_bHasLeftChild) {
currentBestMatch = recursiveSearchNearestNeighbour(nodeIndex + 1, point, distSquared, predicate);
}
// Right side
else if(!isAtLeft && node.m_nRightChildIndex < m_Nodes.size()) {
currentBestMatch = recursiveSearchNearestNeighbour(node.m_nRightChildIndex, point, distSquared, predicate);
}
// Test the current node against the actual best match
if(candidateDistSquared < distSquared && predicate(index)) {
currentBestMatch = nodeIndex;
distSquared = candidateDistSquared;
}
uint32_t candidate = KdNode::NO_NODE;
float axisDistSquared = sqr(point[axis] - node.m_fSplitPosition);
// If the separation axis is closer to the point that the actual best match
// we must search the other side
if(axisDistSquared < distSquared) {
if(isAtLeft && node.m_nRightChildIndex < m_Nodes.size()) {
candidate = recursiveSearchNearestNeighbour(node.m_nRightChildIndex, point, distSquared, predicate);
}
else if(!isAtLeft && node.m_bHasLeftChild) {
candidate = recursiveSearchNearestNeighbour(nodeIndex + 1, point, distSquared, predicate);
}
}
if(candidate != KdNode::NO_NODE) {
currentBestMatch = candidate;
}
return currentBestMatch;
}
void recursiveSearchKNearestNeighbours(uint32_t nodeIndex, const Vec3f& point, size_t K, Predicate predicate,
float& distSquared, std::vector<std::pair<uint32_t, float>>& heap) const {
auto node = m_Nodes[nodeIndex];
uint32_t index = m_NodesData[nodeIndex].m_nIndex;
float candidateDistSquared = sqr_distance(point, m_NodesData[nodeIndex].m_Position);
// If the node is a leaf, end of the recursion
if(node.m_nSplitAxis == 3) {
if((heap.size() < K || candidateDistSquared < distSquared) && predicate(index)) {
heap.push_back(std::make_pair(nodeIndex, candidateDistSquared));
std::push_heap(heap.begin(), heap.end(), compare);
if(heap.size() > K) {
std::pop_heap(heap.begin(), heap.end(), compare);
heap.pop_back();
}
distSquared = heap.front().second;
}
return;
}
uint8_t axis = node.m_nSplitAxis;
bool isAtLeft = (point[axis] <= node.m_fSplitPosition);
// Left side
if(isAtLeft && node.m_bHasLeftChild) {
recursiveSearchKNearestNeighbours(nodeIndex + 1, point, K, predicate, distSquared, heap);
}
// Right side
else if(!isAtLeft && node.m_nRightChildIndex < m_Nodes.size()) {
recursiveSearchKNearestNeighbours(node.m_nRightChildIndex, point, K, predicate, distSquared, heap);
}
// Test the current node against the actual best match
if((heap.size() < K || candidateDistSquared < distSquared) && predicate(index)) {
heap.push_back(std::make_pair(nodeIndex, candidateDistSquared));
std::push_heap(heap.begin(), heap.end(), compare);
if(heap.size() > K) {
std::pop_heap(heap.begin(), heap.end(), compare);
heap.pop_back();
}
distSquared = heap.front().second;
}
float axisDistSquared = sqr(point[axis] - node.m_fSplitPosition);
// If the separation axis is closer to the point than the actual best match
// we must search the other side
if(heap.size() < K || axisDistSquared < distSquared) {
if(isAtLeft && node.m_nRightChildIndex < m_Nodes.size()) {
recursiveSearchKNearestNeighbours(node.m_nRightChildIndex, point, K, predicate, distSquared, heap);
}
else if(!isAtLeft && node.m_bHasLeftChild) {
recursiveSearchKNearestNeighbours(nodeIndex + 1, point, K, predicate, distSquared, heap);
}
}
}
/**
* Returns the distance between this vector and another one.
* @param v The other vector.
*/
T distance(const vector_2d& v) const {
return sqrt (sqr_distance (v));
}
