299 lines
9.9 KiB
C#
299 lines
9.9 KiB
C#
using RobotNet10.GlobalPathPlanner.Model;
|
|
|
|
namespace RobotNet10.GlobalPathPlanner.Space;
|
|
|
|
/// <summary>
|
|
/// KD-Tree implementation for efficient 2D spatial search of nodes.
|
|
/// Thread-safe for read operations after construction.
|
|
/// Time complexity: O(n log n) build, O(log n) search average case.
|
|
/// </summary>
|
|
public class KDTree
|
|
{
|
|
private readonly KDTreeNode? _root;
|
|
private readonly List<GlobalNode> _nodes;
|
|
|
|
/// <summary>
|
|
/// Initializes a new KD-Tree from a collection of nodes.
|
|
/// </summary>
|
|
/// <param name="nodes">The nodes to index. Original collection is not modified.</param>
|
|
/// <exception cref="ArgumentNullException">Thrown when nodes is null.</exception>
|
|
public KDTree(IEnumerable<GlobalNode> nodes)
|
|
{
|
|
ArgumentNullException.ThrowIfNull(nodes);
|
|
|
|
// Create a copy to avoid mutating input
|
|
_nodes = [.. nodes];
|
|
|
|
if (_nodes.Count == 0)
|
|
{
|
|
_root = null;
|
|
return;
|
|
}
|
|
|
|
_root = BuildTree(0, _nodes.Count - 1, 0);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Gets the total number of nodes in the tree.
|
|
/// </summary>
|
|
public int Count => _nodes.Count;
|
|
|
|
/// <summary>
|
|
/// Builds the KD-Tree using index-based recursion to avoid memory allocation overhead.
|
|
/// Time complexity: O(n log n)
|
|
/// </summary>
|
|
private KDTreeNode? BuildTree(int start, int end, int depth)
|
|
{
|
|
if (start > end)
|
|
return null;
|
|
|
|
int axis = depth % 2;
|
|
|
|
// Use QuickSelect to find median without full sort
|
|
int medianIndex = QuickSelect(start, end, (start + end) / 2, axis);
|
|
|
|
return new KDTreeNode(
|
|
node: _nodes[medianIndex],
|
|
axis: axis,
|
|
left: BuildTree(start, medianIndex - 1, depth + 1),
|
|
right: BuildTree(medianIndex + 1, end, depth + 1)
|
|
);
|
|
}
|
|
|
|
/// <summary>
|
|
/// QuickSelect algorithm to find the k-th smallest element.
|
|
/// Time complexity: O(n) average, O(n²) worst case.
|
|
/// </summary>
|
|
private int QuickSelect(int left, int right, int k, int axis)
|
|
{
|
|
while (left < right)
|
|
{
|
|
int pivotIndex = Partition(left, right, axis);
|
|
|
|
if (pivotIndex == k)
|
|
return k;
|
|
else if (k < pivotIndex)
|
|
right = pivotIndex - 1;
|
|
else
|
|
left = pivotIndex + 1;
|
|
}
|
|
return left;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Partitions the array for QuickSelect using median-of-three pivot selection.
|
|
/// This is the CORRECTED version that handles all edge cases properly.
|
|
/// </summary>
|
|
private int Partition(int left, int right, int axis)
|
|
{
|
|
// Handle small subarrays
|
|
if (right - left < 2)
|
|
{
|
|
if (right > left && CompareNodes(_nodes[right], _nodes[left], axis) < 0)
|
|
Swap(left, right);
|
|
return left;
|
|
}
|
|
|
|
// Use median-of-three pivot selection for better performance
|
|
int mid = left + (right - left) / 2; // Overflow-safe
|
|
|
|
// Sort left, mid, right to get median as pivot
|
|
if (CompareNodes(_nodes[mid], _nodes[left], axis) < 0)
|
|
Swap(left, mid);
|
|
if (CompareNodes(_nodes[right], _nodes[left], axis) < 0)
|
|
Swap(left, right);
|
|
if (CompareNodes(_nodes[right], _nodes[mid], axis) < 0)
|
|
Swap(mid, right);
|
|
|
|
// Now: nodes[left] <= nodes[mid] <= nodes[right]
|
|
// Use mid as pivot and hide it at right-1
|
|
GlobalNode pivot = _nodes[mid];
|
|
Swap(mid, right - 1);
|
|
|
|
// Partition with pivot at right-1
|
|
int i = left + 1; // Start after left (which is already <= pivot)
|
|
int j = right - 2; // Start before pivot position
|
|
|
|
while (i <= j)
|
|
{
|
|
// Find element >= pivot from left
|
|
while (i <= j && CompareNodes(_nodes[i], pivot, axis) < 0)
|
|
i++;
|
|
|
|
// Find element <= pivot from right
|
|
while (i <= j && CompareNodes(_nodes[j], pivot, axis) > 0)
|
|
j--;
|
|
|
|
if (i < j)
|
|
{
|
|
Swap(i, j);
|
|
i++;
|
|
j--;
|
|
}
|
|
else
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Put pivot in its final position
|
|
Swap(i, right - 1);
|
|
return i;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Compares two nodes along the specified axis.
|
|
/// </summary>
|
|
private static int CompareNodes(GlobalNode a, GlobalNode b, int axis)
|
|
{
|
|
return axis == 0 ? a.X.CompareTo(b.X) : a.Y.CompareTo(b.Y);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Swaps two elements in the nodes list using tuple deconstruction.
|
|
/// </summary>
|
|
private void Swap(int i, int j)
|
|
{
|
|
if (i != j)
|
|
{
|
|
(_nodes[j], _nodes[i]) = (_nodes[i], _nodes[j]);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Finds the nearest node to the given coordinates within the specified distance limit.
|
|
/// Time complexity: O(log n) average, O(n) worst case.
|
|
/// </summary>
|
|
/// <param name="x">X coordinate</param>
|
|
/// <param name="y">Y coordinate</param>
|
|
/// <param name="limitDistance">Maximum search distance</param>
|
|
/// <returns>The nearest node within limit, or null if none found</returns>
|
|
/// <exception cref="ArgumentOutOfRangeException">Thrown when limitDistance is negative.</exception>
|
|
public GlobalNode? FindNearest(double x, double y, double limitDistance)
|
|
{
|
|
ArgumentOutOfRangeException.ThrowIfNegative(limitDistance);
|
|
|
|
if (_root == null)
|
|
return null;
|
|
|
|
double limitDistSquared = limitDistance * limitDistance;
|
|
var result = FindNearestRecursive(_root, x, y, null, double.MaxValue, limitDistSquared);
|
|
|
|
return result.BestNode;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Recursively finds the nearest node using squared distances to avoid sqrt operations.
|
|
/// </summary>
|
|
private static SearchResult FindNearestRecursive(
|
|
KDTreeNode? node,
|
|
double x,
|
|
double y,
|
|
GlobalNode? bestNode,
|
|
double bestDistSquared,
|
|
double limitDistSquared)
|
|
{
|
|
if (node == null)
|
|
return new SearchResult(bestNode, bestDistSquared);
|
|
|
|
// Calculate squared distance (avoid sqrt for performance)
|
|
double dx = node.Node.X - x;
|
|
double dy = node.Node.Y - y;
|
|
double distSquared = dx * dx + dy * dy;
|
|
|
|
// Update best if this node is closer and within limit
|
|
if (distSquared < bestDistSquared && distSquared <= limitDistSquared)
|
|
{
|
|
bestNode = node.Node;
|
|
bestDistSquared = distSquared;
|
|
}
|
|
|
|
// Determine which side to search first
|
|
double delta = node.Axis == 0 ? x - node.Node.X : y - node.Node.Y;
|
|
KDTreeNode? nearSide = delta < 0 ? node.Left : node.Right;
|
|
KDTreeNode? farSide = delta < 0 ? node.Right : node.Left;
|
|
|
|
// Search near side
|
|
var result = FindNearestRecursive(nearSide, x, y, bestNode, bestDistSquared, limitDistSquared);
|
|
bestNode = result.BestNode;
|
|
bestDistSquared = result.BestDistSquared;
|
|
|
|
// Only search far side if it could contain a closer point
|
|
double deltaSquared = delta * delta;
|
|
if (deltaSquared < bestDistSquared)
|
|
{
|
|
result = FindNearestRecursive(farSide, x, y, bestNode, bestDistSquared, limitDistSquared);
|
|
bestNode = result.BestNode;
|
|
bestDistSquared = result.BestDistSquared;
|
|
}
|
|
|
|
return new SearchResult(bestNode, bestDistSquared);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Finds all nodes within the specified radius from the given coordinates.
|
|
/// </summary>
|
|
/// <param name="x">X coordinate</param>
|
|
/// <param name="y">Y coordinate</param>
|
|
/// <param name="radius">Search radius</param>
|
|
/// <returns>List of all nodes within the radius</returns>
|
|
public List<GlobalNode> FindInRadius(double x, double y, double radius)
|
|
{
|
|
ArgumentOutOfRangeException.ThrowIfNegative(radius);
|
|
|
|
if (_root == null)
|
|
return [];
|
|
|
|
var result = new List<GlobalNode>();
|
|
double radiusSquared = radius * radius;
|
|
FindInRadiusRecursive(_root, x, y, radiusSquared, result);
|
|
return result;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Recursively finds all nodes within radius.
|
|
/// </summary>
|
|
private static void FindInRadiusRecursive(
|
|
KDTreeNode? node,
|
|
double x,
|
|
double y,
|
|
double radiusSquared,
|
|
List<GlobalNode> result)
|
|
{
|
|
if (node == null)
|
|
return;
|
|
|
|
double dx = node.Node.X - x;
|
|
double dy = node.Node.Y - y;
|
|
double distSquared = dx * dx + dy * dy;
|
|
|
|
if (distSquared <= radiusSquared)
|
|
result.Add(node.Node);
|
|
|
|
double delta = node.Axis == 0 ? x - node.Node.X : y - node.Node.Y;
|
|
double deltaSquared = delta * delta;
|
|
|
|
// Search both sides if sphere intersects splitting plane
|
|
if (delta < 0)
|
|
{
|
|
FindInRadiusRecursive(node.Left, x, y, radiusSquared, result);
|
|
if (deltaSquared <= radiusSquared)
|
|
FindInRadiusRecursive(node.Right, x, y, radiusSquared, result);
|
|
}
|
|
else
|
|
{
|
|
FindInRadiusRecursive(node.Right, x, y, radiusSquared, result);
|
|
if (deltaSquared <= radiusSquared)
|
|
FindInRadiusRecursive(node.Left, x, y, radiusSquared, result);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Internal struct to return search results without allocations.
|
|
/// </summary>
|
|
private readonly struct SearchResult(GlobalNode? bestNode, double bestDistSquared)
|
|
{
|
|
public readonly GlobalNode? BestNode = bestNode;
|
|
public readonly double BestDistSquared = bestDistSquared;
|
|
}
|
|
} |