Skip to content

Instantly share code, notes, and snippets.

@wallstop
Created March 15, 2022 01:57
Show Gist options
  • Save wallstop/75b006ad0a6915308e87c955d65e0003 to your computer and use it in GitHub Desktop.
Save wallstop/75b006ad0a6915308e87c955d65e0003 to your computer and use it in GitHub Desktop.
QuadTree
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>>();
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment