QuadTree

QuadTree.cs

namespace Core.DataStructure
{
    using System;
    using System.Collections.Concurrent;
    using System.Collections.Generic;
    using System.Linq;
    using Extension;
    using UnityEngine;

    public sealed class QuadTree<T>
    {
        private sealed class QuadTreeNode<V>
        {
            private static readonly int NumChildren = 4;

            public readonly Bounds boundary;
            public readonly IReadOnlyCollection<QuadTreeNode<V>> children;
            public readonly IReadOnlyCollection<V> elements;
            public readonly bool isTerminal;

            public QuadTreeNode(IReadOnlyCollection<V> elements, Func<V, Vector2> elementTransformer, Bounds boundary,
                int bucketSize)
            {
                this.boundary = boundary;
                isTerminal = elements.Count <= bucketSize;
                List<QuadTreeNode<V>> childrenList = new(isTerminal ? 0 : NumChildren);
                children = childrenList;
                this.elements = elements;
                if (isTerminal)
                {
                    return;
                }

                Vector3 quadrantSize = boundary.size / 2f;
                Vector2 halfQuadrantSize = quadrantSize / 2f;

                Bounds[] quadrants =
                {
                    new Bounds(new Vector3(boundary.center.x - halfQuadrantSize.x, boundary.center.y + halfQuadrantSize.y, boundary.center.z), quadrantSize),
                    new Bounds(new Vector3(boundary.center.x + halfQuadrantSize.x, boundary.center.y + halfQuadrantSize.y, boundary.center.z), quadrantSize),
                    new Bounds(new Vector3(boundary.center.x + halfQuadrantSize.x, boundary.center.y - halfQuadrantSize.y, boundary.center.z), quadrantSize),
                    new Bounds(new Vector3(boundary.center.x - halfQuadrantSize.x, boundary.center.y - halfQuadrantSize.y, boundary.center.z), quadrantSize),
                };

                foreach (Bounds quadrant in quadrants)
                {
                    List<V> pointsInRange = elements
                        .Where(element => quadrant.Contains(elementTransformer(element))).ToList();
                    QuadTreeNode<V> child = new QuadTreeNode<V>(pointsInRange, elementTransformer, quadrant, bucketSize);
                    childrenList.Add(child);
                }
            }
        }

        private const int DefaultBucketSize = 12;

        public readonly IReadOnlyCollection<T> elements;

        private readonly Bounds _bounds;
        private readonly Func<T, Vector2> _elementTransformer;
        private readonly QuadTreeNode<T> _head;
        /*
            Performance optimization - cache the Queues used for GetElementsInBounds.
            These are accessed from the job system, so the stack needs to be thread safe.         
        */
        private readonly ConcurrentStack<Queue<QuadTreeNode<T>>> _nodesToVisit = new();

        public QuadTree(IEnumerable<T> points, Func<T, Vector2> elementTransformer, Bounds boundary,
            int bucketSize = DefaultBucketSize)
        {
            _elementTransformer = elementTransformer;
            _bounds = boundary;
            elements = points.ToList();
            _head = new QuadTreeNode<T>(elements, elementTransformer, boundary, bucketSize);
        }

        public IEnumerable<T> GetElementsInRange(Vector2 position, float range)
        {
            Circle area = new(position, range);
            return GetElementsInBounds(new Bounds(new Vector3(position.x, position.y, 0f),
                    new Vector3(range * 2f, range * 2f, 1f)))
                .Where(element => area.Contains(_elementTransformer(element)));
        }

        public IEnumerable<T> GetElementsInBounds(Bounds bounds)
        {
            if (!bounds.FastIntersects2D(_bounds))
            {
                yield break;
            }

            Queue<QuadTreeNode<T>> nodesToVisit = GetNodesToVisit();
            nodesToVisit.Enqueue(_head);

            do
            {
                QuadTreeNode<T> currentNode = nodesToVisit.Dequeue();
                if (!bounds.FastIntersects2D(currentNode.boundary))
                {
                    continue;
                }
                
                if (currentNode.isTerminal)
                {
                    foreach (T element in currentNode.elements)
                    {
                        if (bounds.Contains(_elementTransformer(element)))
                        {
                            yield return element;
                        }
                    }

                    continue;
                }

                if (bounds.Overlaps2D(currentNode.boundary))
                {
                    foreach (T element in currentNode.elements)
                    {
                        yield return element;
                    }

                    continue;
                }

                foreach (QuadTreeNode<T> child in currentNode.children)
                {
                    nodesToVisit.Enqueue(child);
                }
            } while (0 < nodesToVisit.Count);

            _nodesToVisit.Push(nodesToVisit);
        }

        private Queue<QuadTreeNode<T>> GetNodesToVisit()
        {
            if (_nodesToVisit.TryPop(out Queue<QuadTreeNode<T>> nodesToVisit))
            {
                nodesToVisit.Clear();
                return nodesToVisit;
            }

            return new Queue<QuadTreeNode<T>>();
        }
    }
}