city-generation/Game/Assets/Scripts/Utils/QuadTree.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using TransportGame.Model;

namespace TransportGame.Utils
{
    public class QuadTree<T>: ICollection<T> where T : IPositionable
    {
        private const int Capacity = 8;

        /// <summary>
        /// Gets the subtree in northwest position
        /// </summary>
        public QuadTree<T> NorthWest { get; private set; }

        /// <summary>
        /// Gets the subtree in northeast position
        /// </summary>
        public QuadTree<T> NorthEast { get; private set; }

        /// <summary>
        /// Gets the subtree in southeast position
        /// </summary>
        public QuadTree<T> SouthEast { get; private set; }

        /// <summary>
        /// Gets the subtree in southwest position
        /// </summary>
        public QuadTree<T> SouthWest { get; private set; }

        /// <summary>
        /// Gets the boundary of this quad tree
        /// </summary>
        public Rectangle Boundary { get; private set; }

        private List<T> points = new List<T>(Capacity);

        /// <summary>
        /// Initializes a quad tree using specified boundaries
        /// </summary>
        /// <param name="left">Left</param>
        /// <param name="top">Top</param>
        /// <param name="right">Right</param>
        /// <param name="bottom">Bottom</param>
        public QuadTree(float left, float top, float right, float bottom)
        {
            Boundary = new Rectangle(left, top, right, bottom);
        }

        /// <summary>
        /// Initializes a quad tree using specified boundaries
        /// </summary>
        /// <param name="boundary">Boundaries</param>
        public QuadTree(Rectangle boundary)
        {
            Boundary = boundary;
        }

        private void Subdivide()
        {
            float midx = Boundary.Left + Boundary.Width / 2f;
            float midy = Boundary.Top + Boundary.Height / 2f;

            NorthWest = new QuadTree<T>(Boundary.Left, Boundary.Top, midx, midy);
            NorthEast = new QuadTree<T>(midx, Boundary.Top, Boundary.Right, midy);
            SouthEast = new QuadTree<T>(midx, midy, Boundary.Right, Boundary.Bottom);
            SouthWest = new QuadTree<T>(Boundary.Left, midy, midx, Boundary.Bottom);

            foreach (var point in points)
                Add(point);

            points.Clear();
        }

        private void Merge()
        {
            foreach (var point in NorthWest)
                points.Add(point);

            foreach (var point in NorthEast)
                points.Add(point);

            foreach (var point in SouthEast)
                points.Add(point);

            foreach (var point in SouthWest)
                points.Add(point);

            NorthWest = null;
            NorthEast = null;
            SouthEast = null;
            SouthWest = null;
        }

        /// <summary>
        /// Adds a point in this quad tree
        /// </summary>
        /// <param name="item">Point</param>
        public void Add(T item)
        {
            // Precondition - point must be inside boundaries
            if (!Boundary.Contains(item.Position))
                throw new ArgumentException("Point must be inside boundaries.");

            // Reached capacity, subdivide
            if (NorthWest == null && points.Count >= Capacity)
                Subdivide();

            // Not divided in subtrees
            if (NorthWest == null)
            {
                if (!points.Any(p => p.Position.Equals(item.Position)))
                    points.Add(item);
            }

            // Add in the right subtree
            else
            {
                if (NorthWest.Boundary.Contains(item.Position))
                    NorthWest.Add(item);

                else if (NorthEast.Boundary.Contains(item.Position))
                    NorthEast.Add(item);

                else if (SouthEast.Boundary.Contains(item.Position))
                    SouthEast.Add(item);

                else SouthWest.Add(item);
            }
        }

        /// <summary>
        /// Empties the entire tree
        /// </summary>
        public void Clear()
        {
            NorthWest = null; 
            NorthEast = null;
            SouthEast = null;
            SouthWest = null;
            points.Clear();
        }

        /// <summary>
        /// Tests if specified point is contained in this tree
        /// </summary>
        /// <param name="item">Point</param>
        /// <returns>True if point is contained</returns>
        public bool Contains(T item)
        {
            if (NorthWest == null)
                return points.Contains(item);
            else
                return NorthWest.Contains(item) || NorthEast.Contains(item) || SouthEast.Contains(item) || SouthWest.Contains(item);
        }

        /// <summary>
        /// Copies all points to specified array
        /// </summary>
        /// <param name="array">Array</param>
        /// <param name="arrayIndex">Index where to start copying</param>
        public void CopyTo(T[] array, int arrayIndex)
        {
            foreach (var point in this)
                array[arrayIndex++] = point;
        }

        /// <summary>
        /// Gets the number of points in this quad tree
        /// </summary>
        public int Count
        {
            get 
            {
                return (NorthWest == null) ? points.Count : NorthWest.Count + NorthEast.Count + SouthEast.Count + SouthWest.Count;
            }
        }

        /// <summary>
        /// Returns true if this container is read-only.
        /// </summary>
        public bool IsReadOnly
        {
            get 
            {
                return false;
            }
        }

        /// <summary>
        /// Removes an item from the tree
        /// </summary>
        /// <param name="item">Item to remove</param>
        /// <returns>True if item was removed</returns>
        public bool Remove(T item)
        {
            if (NorthWest == null)
                return points.Remove(item);

            else
            {
                bool result = NorthWest.Remove(item) || NorthEast.Remove(item) || SouthEast.Remove(item) || SouthWest.Remove(item);
                
                // We can merge subdivisions
                if (Count < (Capacity - Capacity / 3))
                    Merge();

                return result;
            }
        }

        /// <summary>
        /// Gets an enumerator for this collection
        /// </summary>
        /// <returns></returns>
        System.Collections.IEnumerator IEnumerable.GetEnumerator()
        {
            foreach (var point in this)
                yield return point;
        }

        /// <summary>
        /// Gets enumerator for this collection
        /// </summary>
        /// <returns></returns>
        public IEnumerator<T> GetEnumerator()
        {
            if (NorthWest == null)
            {
                foreach (var point in points)
                    yield return point;
            }

            else
            {
                foreach (var point in NorthWest)
                    yield return point;

                foreach (var point in NorthEast)
                    yield return point;

                foreach (var point in SouthEast)
                    yield return point;

                foreach (var point in SouthWest)
                    yield return point;
            }
        }

        /// <summary>
        /// Gets all the points in the specified region
        /// </summary>
        /// <param name="rect">Region</param>
        /// <returns>Points</returns>
        public IEnumerable<T> Query(Rectangle rect)
        {
            // No intersection
            if (!Rectangle.Intersects(rect, Boundary))
                return Enumerable.Empty<T>();

            if (NorthWest == null)
            {
                return points.Where(p => rect.Contains(p.Position));
            }
            else
            {
                return NorthWest.Query(rect)
                    .Concat(NorthEast.Query(rect))
                    .Concat(SouthEast.Query(rect))
                    .Concat(SouthWest.Query(rect));
            }
        }
        
        /// <summary>
        /// Gets all the points in the specified region
        /// </summary>
        /// <param name="left">Left</param>
        /// <param name="top">Top</param>
        /// <param name="right">Right</param>
        /// <param name="bottom">Bottom</param>
        /// <returns>Points</returns>
        public IEnumerable<T> Query(float left, float top, float right, float bottom)
        {
            return Query(new Rectangle(left, top, right, bottom));
        }
    }
}