// Taken from https://github.com/sphinxy/DataStructures/blob/35f41dcb29e5ee6d2217fc5ae4525f0990a6be0a/DataStructures/PriorityQueue.cs /* The MIT License (MIT) Copyright (c) 2015 Denis Shulepov Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */ using System; using System.Collections; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; namespace Robust.Shared.Utility { /// /// Heap-based resizable max-priority queue. /// Elements with high priority are served before elements with low priority. /// Priority is defined by comparing elements, so to separate priority from value use /// KeyValuePair or a custom class and provide corresponding Comparer. /// /// Any comparable type, either through a specified Comparer or implementing IComparable<> public sealed class PriorityQueue : ICollection { private readonly IComparer _comparer; private T[] _heap; private const int DEFAULT_CAPACITY = 10; private const int SHRINK_RATIO = 4; private const int RESIZE_FACTOR = 2; private int _shrinkBound; // ReSharper disable once StaticFieldInGenericType private static readonly InvalidOperationException EmptyCollectionException = new("Collection is empty."); /// /// Create a max-priority queue with default capacity of 10. /// /// Custom comparer to compare elements. If omitted - default will be used. public PriorityQueue(IComparer? comparer = null) : this(DEFAULT_CAPACITY, comparer) { } /// /// Create a max-priority queue with provided capacity. /// /// Initial capacity /// Custom comparer to compare elements. If omitted - default will be used. /// Throws when capacity is less than or equal to zero. /// Throws when comparer is null and does not implement IComparable. public PriorityQueue(int capacity, IComparer? comparer = null) { if (capacity <= 0) throw new ArgumentOutOfRangeException("capacity", "Expected capacity greater than zero."); // If no comparer then T must be comparable if (comparer == null && !(typeof(IComparable).IsAssignableFrom(typeof(T)) || typeof(IComparable).IsAssignableFrom(typeof(T)))) { throw new ArgumentException("Expected a comparer for types, which do not implement IComparable.", "comparer"); } _comparer = comparer ?? Comparer.Default; _shrinkBound = capacity / SHRINK_RATIO; _heap = new T[capacity]; } /// /// Current queue capacity /// public int Capacity { get { return _heap.Length; } } public IEnumerator GetEnumerator() { var array = new T[Count]; CopyTo(array, 0); return ((IEnumerable )array).GetEnumerator(); } IEnumerator IEnumerable.GetEnumerator() { return GetEnumerator(); } public void Add(T item) { if (Count == Capacity) GrowCapacity(); _heap[Count++] = item; // provide the index of the last item as for 1-based heap, but also set shift to -1 _heap.Sift(Count, _comparer, shift: -1); // move item "up" until heap principles are not met } /// /// Removes and returns a max element from the priority queue. /// /// Max element in the collection /// Throws when queue is empty. public T Take() { if (Count == 0) throw EmptyCollectionException; var item = _heap[0]; Count--; _heap.Swap(0, Count); // last element at count _heap[Count] = default(T)!; // release hold on the object // provide index of first item as for 1-based heap, but also set shift to -1 _heap.Sink(1, Count, _comparer, shift: -1); // move item "down" while heap principles are not met if (Count <= _shrinkBound && Count > DEFAULT_CAPACITY) { ShrinkCapacity(); } return item; } /// /// Returns a max element from the priority queue without removing it. /// /// Max element in the collection /// Throws when queue is empty. public T Peek() { if (Count == 0) throw EmptyCollectionException; return _heap[0]; } public void Clear() { _heap = new T[DEFAULT_CAPACITY]; Count = 0; } public bool Contains(T item) { return GetItemIndex(item) >= 0; } public void CopyTo(T[] array, int arrayIndex) { if (array == null) throw new ArgumentNullException("array"); if (arrayIndex < 0) throw new ArgumentOutOfRangeException("arrayIndex"); if (array.Length - arrayIndex < Count) throw new ArgumentException("Insufficient space in destination array."); Array.Copy(_heap, 0, array, arrayIndex, Count); array.HeapSort(arrayIndex, Count, _comparer); } public bool Remove(T item) { var index = GetItemIndex(item); switch (index) { case -1: return false; case 0: Take(); break; default: // provide a 1-based index of the item RemoveAt(index + 1, shift: -1); break; } return true; } public int Count { get; private set; } public bool IsReadOnly { get { return false; } } /// /// Removes item at given index /// /// 1-based index of the element to remove /// Shift allows to compensate and work with arrays where heap starts not from the element at position 1. /// Shift -1 allows to work with 0-based heap as if it was 1-based. But the main reason for this is the CopyTo method. private void RemoveAt(int index, int shift) { var itemIndex = index + shift; Count--; _heap.Swap(itemIndex, Count); // last element at Count _heap[Count] = default(T)!; // release hold on the object // use a 1-based-heap index and then apply shift of -1 var parent = index / 2 + shift; // get parent // if new item at index is greater than it's parent then sift it up, else sink it down if (_comparer.GreaterOrEqual(_heap[itemIndex], _heap[parent])) { // provide a 1-based-heap index _heap.Sift(index, _comparer, shift); } else { // provide a 1-based-heap index _heap.Sink(index, Count, _comparer, shift); } } /// /// Returns the real index of the first occurrence of the given item or -1. /// private int GetItemIndex(T item) { for (int i = 0; i < Count; i++) { if (_comparer.Compare(_heap[i], item) == 0) return i; } return -1; } private void GrowCapacity() { int newCapacity = Capacity * RESIZE_FACTOR; Array.Resize(ref _heap, newCapacity); // first element is at position 1 _shrinkBound = newCapacity / SHRINK_RATIO; } private void ShrinkCapacity() { int newCapacity = Capacity / RESIZE_FACTOR; Array.Resize(ref _heap, newCapacity); // first element is at position 1 _shrinkBound = newCapacity / SHRINK_RATIO; } } internal static class HeapMethods { internal static void Swap(this T[] array, int i, int j) { var tmp = array[i]; array[i] = array[j]; array[j] = tmp; } internal static bool GreaterOrEqual(this IComparer comparer, [AllowNull] T x, [AllowNull] T y) { return comparer.Compare(x, y) >= 0; } /// /// Moves the item with given index "down" the heap while heap principles are not met. /// /// Any comparable type /// Array, containing the heap /// 1-based index of the element to sink /// Number of items in the heap /// Comparer to compare the items /// Shift allows to compensate and work with arrays where heap starts not from the element at position 1. /// Shift -1 allows to work with 0-based heap as if it was 1-based. But the main reason for this is the CopyTo method. /// internal static void Sink(this T[] heap, int i, int count, IComparer comparer, int shift) { var lastIndex = count + shift; while (true) { var itemIndex = i + shift; var leftIndex = 2 * i + shift; if (leftIndex > lastIndex) { return; // reached last item } var rightIndex = leftIndex + 1; var hasRight = rightIndex <= lastIndex; var item = heap[itemIndex]; var left = heap[leftIndex]; var right = hasRight ? heap[rightIndex] : default(T); // if item is greater than children - heap is fine, exit // ReSharper disable once RedundantTypeArgumentsOfMethod if (GreaterOrEqual(comparer, item, left) && (!hasRight || GreaterOrEqual(comparer, item, right))) { return; } // else exchange with greater of children // ReSharper disable once RedundantTypeArgumentsOfMethod var greaterChildIndex = !hasRight || GreaterOrEqual(comparer, left, right) ? leftIndex : rightIndex; heap.Swap(itemIndex, greaterChildIndex); // continue at new position i = greaterChildIndex - shift; } } /// /// Moves the item with given index "up" the heap while heap principles are not met. /// /// Any comparable type /// Array, containing the heap /// 1-based index of the element to sink /// Comparer to compare the items /// Shift allows to compensate and work with arrays where heap starts not from the element at position 1. /// Value -1 allows to work with 0-based heap as if it was 1-based. But the main reason for this is the CopyTo method. /// internal static void Sift(this T[] heap, int i, IComparer comparer, int shift) { while (true) { if (i <= 1) return; // reached root var parent = i / 2 + shift; // get parent var index = i + shift; // if root is greater or equal - exit if (GreaterOrEqual(comparer, heap[parent], heap[index])) { return; } heap.Swap(parent, index); i = parent - shift; } } /// /// Sorts the heap in descending order. /// /// Any comparable type /// Array, containing the heap /// Index in the array, from which the heap structure begins /// Number of items in the heap /// Comparer to compare the items internal static void HeapSort(this T[] heap, int startIndex, int count, IComparer comparer) { var shift = startIndex - 1; var lastIndex = startIndex + count; var left = count; // take the max item and exchange it with the last one // decrement the count of items in the heap and sink the first item to restore the heap rules // repeat for every element in the heap while (lastIndex > startIndex) { lastIndex--; left--; heap.Swap(startIndex, lastIndex); heap.Sink(1, left, comparer, shift); } // when done items will be sorted in ascending order, but this is a Max-PriorityQueue, so reverse Array.Reverse(heap, startIndex, count); } } }