Unity的二叉堆优先队列

using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace GPWiki
{
    /// <summary>
    /// 改进后的二叉堆（最小堆），支持 decrease-key 操作
    /// </summary>
    public class BinaryHeap<T> : ICollection<T> where T : IComparable<T>
    {
        private const int DEFAULT_SIZE = 4;
        private T[] _data;
        private int _count;
        private int _capacity;
        // 用于记录每个元素在堆中的索引
        private Dictionary<T, int> indexMap;

        public int Count { get { return _count; } }

        public int Capacity
        {
            get { return _capacity; }
            set
            {
                _capacity = Math.Max(value, _count);
                if (_capacity != _data.Length)
                {
                    T[] newData = new T[_capacity];
                    Array.Copy(_data, newData, _count);
                    _data = newData;
                }
            }
        }

        public BinaryHeap()
        {
            _capacity = DEFAULT_SIZE;
            _data = new T[_capacity];
            _count = 0;
            indexMap = new Dictionary<T, int>();
        }

        public T Peek()
        {
            if (_count == 0)
                throw new InvalidOperationException("Heap is empty.");
            return _data[0];
        }

        public void Add(T item)
        {
            if (_count == _capacity)
                Capacity = _capacity * 2;
            _data[_count] = item;
            indexMap[item] = _count;
            UpHeap(_count);
            _count++;
        }

        public T Remove()
        {
            if (_count == 0)
                throw new InvalidOperationException("Heap is empty.");
            T root = _data[0];
            Swap(0, _count - 1);
            _count--;
            indexMap.Remove(root);
            DownHeap(0);
            return root;
        }

        /// <summary>
        /// 降低某个元素的关键值（这里通过移除旧记录再添加新记录实现）
        /// </summary>
        public void DecreaseKey(T item, T newItem)
        {
            if (!indexMap.TryGetValue(item, out int index))
                return;
            _data[index] = newItem;
            indexMap.Remove(item);
            indexMap[newItem] = index;
            UpHeap(index);
        }

        /// <summary>
        /// 交换数组中两个元素，并同步更新 indexMap
        /// </summary>
        private void Swap(int i, int j)
        {
            T temp = _data[i];
            _data[i] = _data[j];
            _data[j] = temp;
            indexMap[_data[i]] = i;
            indexMap[_data[j]] = j;
        }

        private void UpHeap(int index)
        {
            while (index > 0)
            {
                int parent = Parent(index);
                if (_data[index].CompareTo(_data[parent]) < 0)
                {
                    Swap(index, parent);
                    index = parent;
                }
                else
                {
                    break;
                }
            }
        }

        private void DownHeap(int index)
        {
            while (true)
            {
                int left = Child1(index);
                int right = Child2(index);
                int smallest = index;
                if (left < _count && _data[left].CompareTo(_data[smallest]) < 0)
                    smallest = left;
                if (right < _count && _data[right].CompareTo(_data[smallest]) < 0)
                    smallest = right;
                if (smallest != index)
                {
                    Swap(index, smallest);
                    index = smallest;
                }
                else
                {
                    break;
                }
            }
        }

        private int Parent(int index) => (index - 1) / 2;
        private int Child1(int index) => index * 2 + 1;
        private int Child2(int index) => index * 2 + 2;

        #region ICollection<T> 实现

        public bool IsReadOnly => false;

        public void Clear()
        {
            _count = 0;
            _data = new T[_capacity];
            indexMap.Clear();
        }

        public bool Contains(T item)
        {
            return indexMap.ContainsKey(item);
        }

        public void CopyTo(T[] array, int arrayIndex)
        {
            Array.Copy(_data, 0, array, arrayIndex, _count);
        }

        public IEnumerator<T> GetEnumerator()
        {
            for (int i = 0; i < _count; i++)
                yield return _data[i];
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        /// <summary>
        /// 删除指定元素（用于更新时删除旧记录）
        /// </summary>
        public bool Remove(T item)
        {
            if (!indexMap.TryGetValue(item, out int index))
                return false;
            Swap(index, _count - 1);
            _count--;
            indexMap.Remove(item);
            DownHeap(index);
            return true;
        }

        #endregion
    }
}