Spatial Partition

using System.Collections.Generic;
using UnityEngine;

namespace Patterns.Optimization.SpatialPartition
{
    /// <summary>
    /// Spawns wandering units, then every frame finds those within a radius
    /// of this object using the grid instead of scanning every unit.
    /// </summary>
    public class SpatialTester : MonoBehaviour
    {
        [SerializeField] private Unit unitPrefab;
        [SerializeField] private int unitCount = 500;
        [SerializeField] private float worldSize = 100f;
        [SerializeField] private float cellSize = 10f;
        [SerializeField] private float queryRadius = 8f;

        private SpatialGrid<Unit> grid;
        private readonly List<Unit> results = new();

        private void Awake()
        {
            grid = new SpatialGrid<Unit>(cellSize);

            for (int i = 0; i < unitCount; i++)
            {
                Vector2 p = new(Random.value * worldSize, Random.value * worldSize);
                Unit unit = Instantiate(unitPrefab, new Vector3(p.x, 0f, p.y), Quaternion.identity);
                unit.Init(grid, worldSize);
            }
        }

        private void Update()
        {
            Vector2 center = new(transform.position.x, transform.position.z);
            int checks = grid.Query(center, queryRadius, results);

            foreach (Unit unit in results)
                Debug.DrawLine(transform.position, unit.transform.position, Color.cyan);

            if (Time.frameCount % 60 == 0)
                Debug.Log($"Found {results.Count} with {checks} distance checks (brute force: {unitCount}).");
        }
    }
}

using System.Collections.Generic;
using UnityEngine;

namespace Patterns.Optimization.SpatialPartition
{
    /// <summary>
    /// A uniform grid that buckets items by cell. Queries only look at the
    /// cells overlapping the search circle.
    /// </summary>
    public class SpatialGrid<T> where T : class, ISpatialItem
    {
        private readonly float cellSize;
        private readonly Dictionary<Vector2Int, List<T>> cells = new();

        public SpatialGrid(float cellSize) => this.cellSize = cellSize;

        public Vector2Int CellOf(Vector2 position) =>
            new(Mathf.FloorToInt(position.x / cellSize), Mathf.FloorToInt(position.y / cellSize));

        public void Add(T item, Vector2Int cell)
        {
            if (!cells.TryGetValue(cell, out var list))
                cells[cell] = list = new List<T>();
            list.Add(item);
        }

        public void Remove(T item, Vector2Int cell)
        {
            if (cells.TryGetValue(cell, out var list))
                list.Remove(item);
        }

        public void Move(T item, Vector2Int from, Vector2Int to)
        {
            if (from == to) return;
            Remove(item, from);
            Add(item, to);
        }

        /// <summary>Fills results; returns how many distance checks ran.</summary>
        public int Query(Vector2 center, float radius, List<T> results)
        {
            results.Clear();
            int checks = 0;
            float radiusSq = radius * radius;
            Vector2Int min = CellOf(center - Vector2.one * radius);
            Vector2Int max = CellOf(center + Vector2.one * radius);

            for (int x = min.x; x <= max.x; x++)
            for (int y = min.y; y <= max.y; y++)
            {
                if (!cells.TryGetValue(new Vector2Int(x, y), out var list)) continue;
                foreach (T item in list)
                {
                    checks++;
                    if ((item.Position - center).sqrMagnitude <= radiusSq)
                        results.Add(item);
                }
            }
            return checks;
        }
    }

    /// <summary>Anything the grid can store.</summary>
    public interface ISpatialItem
    {
        Vector2 Position { get; }
    }
}

using UnityEngine;

namespace Patterns.Optimization.SpatialPartition
{
    /// <summary>
    /// A wandering unit that keeps the grid informed whenever it crosses
    /// into a new cell.
    /// </summary>
    public class Unit : MonoBehaviour, ISpatialItem
    {
        [SerializeField] private float speed = 4f;

        private SpatialGrid<Unit> grid;
        private Vector2Int cell;
        private Vector2 velocity;
        private float bounds;

        public Vector2 Position => new(transform.position.x, transform.position.z);

        public void Init(SpatialGrid<Unit> owner, float worldSize)
        {
            grid = owner;
            bounds = worldSize;
            velocity = Random.insideUnitCircle.normalized * speed;
            cell = grid.CellOf(Position);
            grid.Add(this, cell);
        }

        private void Update()
        {
            Vector2 p = Position + velocity * Time.deltaTime;
            if (p.x < 0f || p.x > bounds) velocity.x = -velocity.x;
            if (p.y < 0f || p.y > bounds) velocity.y = -velocity.y;
            p = new Vector2(Mathf.Clamp(p.x, 0f, bounds), Mathf.Clamp(p.y, 0f, bounds));
            transform.position = new Vector3(p.x, 0f, p.y);

            // Re-bucket only when the cell actually changes.
            Vector2Int newCell = grid.CellOf(p);
            if (newCell != cell)
            {
                grid.Move(this, cell, newCell);
                cell = newCell;
            }
        }

        private void OnDestroy() => grid?.Remove(this, cell);
    }
}