PriorityQueue.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See LICENSE in the project root for license information.

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace HoloToolkit.Unity
{
    /// <summary>
    /// Min-heap priority queue. In other words, lower priorities will be removed from the queue first.
    /// See http://en.wikipedia.org/wiki/Binary_heap for more info.
    /// </summary>
    /// <typeparam name="TPriority">Type for the priority used for ordering.</typeparam>
    /// <typeparam name="TValue">Type of values in the queue.</typeparam>
    class PriorityQueue<TPriority, TValue> : IEnumerable<KeyValuePair<TPriority, TValue>>
    {
        public class ValueCollection : IEnumerable<TValue>
        {
            private readonly PriorityQueue<TPriority, TValue> parentCollection;

            public ValueCollection(PriorityQueue<TPriority, TValue> parentCollection)
            {
                this.parentCollection = parentCollection;
            }

            #region IEnumerable

            public IEnumerator<TValue> GetEnumerator()
            {
                foreach (KeyValuePair<TPriority, TValue> pair in parentCollection)
                {
                    yield return pair.Value;
                }
            }

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

            #endregion
        }

        private readonly IComparer<TPriority> priorityComparer;

        public PriorityQueue() : this(Comparer<TPriority>.Default) { }

        public PriorityQueue(IComparer<TPriority> comparer)
        {
            if (comparer == null)
            {
                throw new ArgumentNullException();
            }

            priorityComparer = comparer;
        }

        private readonly List<KeyValuePair<TPriority, TValue>> queue = new List<KeyValuePair<TPriority, TValue>>();
        private ValueCollection valueCollection;

        public ValueCollection Values
        {
            get
            {
                if (valueCollection == null)
                {
                    valueCollection = new ValueCollection(this);
                }

                return valueCollection;
            }
        }

        #region IEnumerable

        public IEnumerator<KeyValuePair<TPriority, TValue>> GetEnumerator()
        {
            return queue.GetEnumerator();
        }

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

        #endregion

        /// <summary>
        /// Clears the priority queue
        /// </summary>
        public void Clear()
        {
            queue.Clear();
        }

        /// <summary>
        /// Add an element to the priority queue.
        /// </summary>
        /// <param name="priority">Priority of the element</param>
        /// <param name="value"></param>
        public void Push(TPriority priority, TValue value)
        {
            queue.Add(new KeyValuePair<TPriority, TValue>(priority, value));
            BubbleUp();
        }

        /// <summary>
        /// Number of elements in priority queue
        /// </summary>
        public int Count
        {
            get
            {
                return queue.Count;
            }
        }

        /// <summary>
        /// Get the element with the minimum priority in the queue. The Key in the return value is the priority.
        /// </summary>
        public KeyValuePair<TPriority, TValue> Top
        {
            get
            {
                return queue[0];
            }
        }

        /// <summary>
        /// Pop the minimal element of the queue. Will fail at runtime if queue is empty.
        /// </summary>
        /// <returns>The minmal element</returns>
        public KeyValuePair<TPriority, TValue> Pop()
        {
            KeyValuePair<TPriority, TValue> ret = queue[0];
            queue[0] = queue[queue.Count - 1];
            queue.RemoveAt(queue.Count - 1);
            BubbleDown();
            return ret;
        }

        /// <summary>
        /// Returns whether or not the value is contained in the queue
        /// </summary>
        public bool Contains(TValue value)
        {
            return queue.Any(itm => EqualityComparer<TValue>.Default.Equals(itm.Value, value));
        }

        /// <summary>
        /// Removes the first element that equals the value from the queue
        /// </summary>
        public bool Remove(TValue value)
        {
            int idx = queue.FindIndex(itm => EqualityComparer<TValue>.Default.Equals(itm.Value, value));
            if (idx == -1)
            {
                return false;
            }

            queue[idx] = queue[queue.Count - 1];
            queue.RemoveAt(queue.Count - 1);
            BubbleDown();

            return true;
        }

        /// <summary>
        /// Removes all elements with this priority from the queue.
        /// </summary>
        /// <returns>True if elements were removed</returns>
        public bool RemoveAtPriority(TPriority priority, Predicate<TValue> shouldRemove)
        {
            bool removed = false;

            for (int i = queue.Count - 1; i >= 0; --i)
            {
                // TODO: early out if key < priority
                if (queue[i].Key.Equals(priority) && (shouldRemove == null || shouldRemove(queue[i].Value)))
                {
                    queue[i] = queue[queue.Count - 1];
                    queue.RemoveAt(queue.Count - 1);
                    BubbleDown();

                    removed = true;
                }
            }

            return removed;
        }

        /// <summary>
        /// Bubble up the last element in the queue until it's in the correct spot.
        /// </summary>
        private void BubbleUp()
        {
            int node = queue.Count - 1;
            while (node > 0)
            {
                int parent = (node - 1) >> 1;
                if (priorityComparer.Compare(queue[parent].Key, queue[node].Key) < 0)
                {
                    break; // we're in the right order, so we're done
                }
                KeyValuePair<TPriority, TValue> tmp = queue[parent];
                queue[parent] = queue[node];
                queue[node] = tmp;
                node = parent;
            }
        }

        /// <summary>
        /// Bubble down the first element until it's in the correct spot.
        /// </summary>
        private void BubbleDown()
        {
            int node = 0;
            while (true)
            {
                // Find smallest child
                int child0 = (node << 1) + 1;
                int child1 = (node << 1) + 2;
                int smallest;
                if (child0 < queue.Count && child1 < queue.Count)
                {
                    smallest = priorityComparer.Compare(queue[child0].Key, queue[child1].Key) < 0 ? child0 : child1;
                }
                else if (child0 < queue.Count)
                {
                    smallest = child0;
                }
                else if (child1 < queue.Count)
                {
                    smallest = child1;
                }
                else
                {
                    break; // 'node' is a leaf, since both children are outside the array
                }

                if (priorityComparer.Compare(queue[node].Key, queue[smallest].Key) < 0)
                {
                    break; // we're in the right order, so we're done.
                }

                KeyValuePair<TPriority, TValue> tmp = queue[node];
                queue[node] = queue[smallest];
                queue[smallest] = tmp;
                node = smallest;
            }
        }
    }
}