| | | 1 | | using System; |
| | | 2 | | using System.Collections.Generic; |
| | | 3 | | using System.Linq; |
| | | 4 | | using UnityEngine; |
| | | 5 | | |
| | | 6 | | // K-D tree node. |
| | | 7 | | public class KDNode<T> { |
| | | 8 | | public T Data; |
| | | 9 | | public KDNode<T> Left = null; |
| | | 10 | | public KDNode<T> Right = null; |
| | | 11 | | |
| | | 12 | | public KDNode(T data) { |
| | | 13 | | Data = data; |
| | | 14 | | } |
| | | 15 | | } |
| | | 16 | | |
| | | 17 | | // K-D tree with 2D coordinates. |
| | | 18 | | public class KDTree<T> { |
| | | 19 | | private KDNode<T> _root; |
| | | 20 | | private Func<T, Vector2> _getCoordinates; |
| | | 21 | | |
| | 0 | 22 | | public KDTree(in IReadOnlyList<T> points, Func<T, Vector2> getCoordinates) { |
| | 0 | 23 | | _getCoordinates = getCoordinates; |
| | 0 | 24 | | _root = BuildTree(points, depth: 0); |
| | 0 | 25 | | } |
| | | 26 | | |
| | 0 | 27 | | private KDNode<T> BuildTree(in IReadOnlyList<T> points, int depth) { |
| | 0 | 28 | | if (points.Count == 0) |
| | 0 | 29 | | return null; |
| | | 30 | | |
| | 0 | 31 | | int k = 2; |
| | 0 | 32 | | int axis = depth % k; |
| | | 33 | | |
| | | 34 | | // Sort the points by axis and find the median. |
| | 0 | 35 | | List<T> sortedPoints = |
| | 0 | 36 | | points.OrderBy(point => (axis == 0 ? _getCoordinates(point).x : _getCoordinates(point).y)) |
| | | 37 | | .ToList(); |
| | 0 | 38 | | int medianIndex = sortedPoints.Count / 2; |
| | 0 | 39 | | T medianPoint = sortedPoints[medianIndex]; |
| | | 40 | | |
| | 0 | 41 | | KDNode<T> node = new KDNode<T>(medianPoint); |
| | 0 | 42 | | List<T> leftPoints = sortedPoints.GetRange(0, medianIndex); |
| | 0 | 43 | | List<T> rightPoints = |
| | | 44 | | sortedPoints.GetRange(medianIndex + 1, sortedPoints.Count - medianIndex - 1); |
| | | 45 | | |
| | 0 | 46 | | node.Left = BuildTree(leftPoints, depth + 1); |
| | 0 | 47 | | node.Right = BuildTree(rightPoints, depth + 1); |
| | 0 | 48 | | return node; |
| | 0 | 49 | | } |
| | | 50 | | |
| | 0 | 51 | | public T NearestNeighbor(Vector2 target) { |
| | 0 | 52 | | KDNode<T> neighbor = NearestNeighbor(_root, target, depth: 0, best: null); |
| | 0 | 53 | | if (neighbor == null) { |
| | 0 | 54 | | return default(T); |
| | | 55 | | } |
| | 0 | 56 | | return neighbor.Data; |
| | 0 | 57 | | } |
| | | 58 | | |
| | 0 | 59 | | private KDNode<T> NearestNeighbor(KDNode<T> node, Vector2 target, int depth, KDNode<T> best) { |
| | 0 | 60 | | if (node == null) |
| | 0 | 61 | | return best; |
| | | 62 | | |
| | 0 | 63 | | float currentDistance = Vector2.Distance(_getCoordinates(node.Data), target); |
| | 0 | 64 | | float bestDistance = |
| | | 65 | | best == null ? float.MaxValue : Vector2.Distance(_getCoordinates(best.Data), target); |
| | 0 | 66 | | if (currentDistance < bestDistance) |
| | 0 | 67 | | best = node; |
| | | 68 | | |
| | 0 | 69 | | int axis = depth % 2; |
| | 0 | 70 | | KDNode<T> nextBranch = (axis == 0 ? target.x < _getCoordinates(node.Data).x |
| | | 71 | | : target.y < _getCoordinates(node.Data).y) |
| | | 72 | | ? node.Left |
| | | 73 | | : node.Right; |
| | 0 | 74 | | KDNode<T> otherBranch = (nextBranch == node.Left) ? node.Right : node.Left; |
| | | 75 | | |
| | | 76 | | // Check the next branch. |
| | 0 | 77 | | best = NearestNeighbor(nextBranch, target, depth + 1, best); |
| | | 78 | | |
| | | 79 | | // Check the other branch. |
| | 0 | 80 | | if ((axis == 0 && Mathf.Abs(target.x - _getCoordinates(node.Data).x) < bestDistance) || |
| | 0 | 81 | | (axis == 1 && Mathf.Abs(target.y - _getCoordinates(node.Data).y) < bestDistance)) { |
| | 0 | 82 | | best = NearestNeighbor(otherBranch, target, depth + 1, best); |
| | 0 | 83 | | } |
| | | 84 | | |
| | 0 | 85 | | return best; |
| | 0 | 86 | | } |
| | | 87 | | } |