.NET 6: PriorityQueue

В .NET 6 появилась новая коллекция — PriorityQueue<TElement,TPriority>. До этого очереди с приоритетами уже были в .NET, но только в виде внутренних классов — они использовались под капотом разных механизмов в WPF, Rx.NET и в других частях фреймворка. 

Но в .NET 6 PriorityQueue стала новой коллекцией, которой теперь можно пользоваться из клиентского кода. Давайте посмотрим, что предлагает эта очередь, как она устроена внутри и насколько быстро работает. Под катом будет постепенное погружение: от примеров использования в коде к введению n-арные деревья.

Что это такое

PriorityQueue — это коллекция элементов, каждый из которых имеет значение и приоритет. Элементы добавляются в очередь с определенным приоритетом и извлекаться из неё “последовательно” — вначале будут извлечены элементы с наименьшим приоритетом, добавленные первыми. То есть, PriorityQueue — это модификация обычной FIFO (first-in-first-out) очереди, в которой элементы с меньшим приоритетом будут оказываться перед всеми элементами с большим приоритетом и извлекаться первыми. 

В каких задачах может пригодиться PriorityQueue? Если коротко, то в тех же, что и обычная очередь, когда часть элементов нужно обработать “вне очереди”. С большой вероятностью такой сценарий понадобится как часть инфраструктурного кода. Вот пара примеров использования PriorityQueue:

1. Организация очереди обработки запросов веб-сервером, в которой входящие запросы из Lisnter’а складываются в очередь и приоритезириуются на основе оставшегося лимита времени запроса или явного приоритета запроса (запрос демона или запрос фронта).

2. Алгоритм Дейкстры — если пишите логистический сервис или просто ищите путь в графе обработчиков в вашем workflow (а это может быть и вывод правил в экспертных системах, и парсинг сайтов/данных). 

3. Любая абстрактная система асинхронной обработки данных/событий/запросов, в которой может быть класс обслуживания — например, если вы хотите дать возможность выполнить в своей системе какой-то служебный запрос перед всеми на случай забитой очереди.

Как работает PriorityQueue

Обычная очередь извлекает элементы в порядке их добавления — first in first out. Очередь с приоритетами извлекает элементы в порядке приоритета от меньшего к большему и также реализует first in first out для элементов с одинаковым приоритетом:

var queue = new PriorityQueue<string, int>();
queue.Enqueue("Миша", 30);
queue.Enqueue("Маша", 29);
queue.Enqueue("Дима", 31);
queue.Enqueue("Коля", 30);
queue.Enqueue("Паша", 31);
 
while (queue.TryDequeue(out var name, out var age))
{
    Console.WriteLine($"{name}, age {age}");
}

Вывод на консоль:

Маша, age 29
Миша, age 30
Коля, age 30
Дима, age 31
Паша, age 31

Как PriorityQueue сравнивает приоритеты

Чтобы определить наименьший приоритет нужно как-то сравнивать приоритеты между собой. PriorityQueue использует для этого стандартный интерфейс IComparer<T>:

public interface IComparer<in T>
{
    int Compare(T? x, T? y);
}

По умолчанию используется дефолтный компоратор System.Collections.Generic.Comparer<TPriority>.Default. Если приоритет задается сложным типом (например, классом ServerHealthState), то можно передать в конструктор PriorityQueue свою реализацию компоратора:

// ServerHealthStateComparer реализует IComparer<ServerHealthState>
var comparer = new ServerHealthStateComparer(); 
var queue = new PriorityQueue<ServerInfo, ServerHealthState>(comparer);

Контракт PriorityQueue

Создать новый экземпляр PriorityQueue можно тремя способами: инициализировать пустую очередь, задать начальное количество элементов чтобы избежать дополнительного расширения массива при добавлении новых элементов или инициализировать очередь готовым набором значений с приоритетом:

public PriorityQueue<TElement,TPriority>()
public PriorityQueue<TElement,TPriority>(int initialCapacity)
public PriorityQueue<TElement,TPriority>(IEnumerable<ValueTuple<TElement,TPriority>> items)

Каждый конструктор имеет перегрузку с компаратором в качестве дополнительного элемента:

public PriorityQueue<TElement,TPriority>(IComparer<TPriority>? comparer)
public PriorityQueue<TElement,TPriority>(int initialCapacity, IComparer<TPriority>? comparer)
public PriorityQueue<TElement,TPriority>(IEnumerable<ValueTuple<TElement,TPriority>> items, IComparer<TPriority> comparer)

Интересной особенностью очереди с приоритетом является то, что в случае пустой очереди начальное хранилище оказывается пустым и расширяется до 4 элементов при первом добавлении:

_nodes = Array.Empty<(TElement, TPriority)>();

У PriorityQueue есть несколько свойств для получения информации об очереди:

// Получает компоратор приоритетов, используемый в текущем экземпляре PriorityQueue
public IComparer<TPriority> Comparer { get; } 

// Получает количество элементов в очереди
public int Count { get; } 

// Получает коллекцию, которая перечисляет элементы очереди в неупорядоченном виде.
public PriorityQueue<TElement,TPriority>.UnorderedItemsCollection UnorderedItems { get; }

Новые элементы можно добавить либо по одному, либо сразу множеством — с общим приоритетом или отдельным значением приоритета для каждого элемента:

// Добавляет в очередь элемент с заданным приоритетом
public void Enqueue(TElement element, TPriority priority);
 
// Добавляет в очередь последовательность элементов с заданным приоритетом
public void EnqueueRange(System.Collections.Generic.IEnumerable<TElement> elements, TPriority priority);
 
// Добавляет в очередь последовательность элементов с заданными приоритетами
public void EnqueueRange(System.Collections.Generic.IEnumerable<(TElement Element, TPriority Priority)> items);

Кроме извлечения очередного элемента из очереди есть привычный метод Peek, который получает элемент из очереди, не удаляя его. В случае с пустой очередью попытка извлечения элемента закончится исключением InvalidOperationException, для безопасного извлечения есть аналогичные методы с префиксом Try:

// Удаляет из очереди и возвращает элемент с наименьшим приоритетом
public TElement Dequeue();
 
// Пытается удалить из очереди и возвращает элемент с наименьшим приоритетом. 
// Возвращает результат операции: true при успешном удалении и false если очередь пуста
public bool TryDequeue(out TElement element, out TPriority priority);
 
// Возвращает элемент с наименьшим приоритетом, не удаляя его из очереди
public TElement Peek();
 
// Пытается вернуть элемент с наименьшим приоритетом, не удаляя его из очереди
// Возвращает результат операции: true при успешном извлечении и false если очередь пуста
public bool TryPeek(out TElement element, out TPriority priority);

Можно быстро очистить очередь при помощи метода Clear:

// Удаляет все элементы из очереди 
public void Clear();

Есть и несколько специализированных методов, которые можно использовать для увеличения эффективности работы с очередью — быстро расширить внутреннее хранилище до определенного значения, обрезать внутреннее хранилище или одновременно извлечь элемент и добавить новый:

// Добавляет в очередь элемент с заданным приоритетом и извлекает элемент с наименьшим приоритетом
public TElement EnqueueDequeue(TElement element, TPriority priority);
    
// Гарантирует, что очередь может хранить capacity элементов без дополнительного расширения внутреннего хранилища
// Расширяет текущее внутреннее хранилище до размера capacity, если его ёмкость была меньше
// Возвращает capacity внутреннего хранилища
public int EnsureCapacity(int capacity); 
 
// Устанавливает capacity внутренннего хранилища на фактическое количество элементов в очереди, если это меньше 90% от текущей capacity 
public void TrimExcess ();

Что внутри PriorityQueue

Как устроена очередь с приоритетами внутри? 

Самая простая реализация “в лоб”, которую можно придумать — это связный список элементов. Новые элементы добавляются в хвост списка за константное время, а при извлечении нужно найти первый элемент с наименьшим приоритетом, перебрав для этого все элементы за линейное время. Peek тоже будет линейным. 

Можно ли сделать извлечение элементов быстрее, чем O(n)? Для этого понадобится держать список в отсортированном упорядоченном состоянии. Если мы используем в качестве структуры данных связный список, то для добавления понадобится последовательно перебирать элементы в поиске нужного места для вставки и асимптотика Enqueue будет O(n), так что теперь добавление в очередь становится слишком дорогостоящей операцией. Если заменять список на динамический массив, то вставка потребует “развигания” элементов и всё равно останется O(n)-операцией.

На самом деле, для решения задачи хранения элементов с приоритетом, “быстрого” добавления новых элементов и извлечения первого элемента с наименьшим приоритетом идеально подходит такая структура данных, как куча. Есть несколько типов универсальных куч, а ещё специализированные кучи, которые предоставляют дополнительную функциональность или работают эффективнее, но имеют ограничения. Например, в качестве типа приоритета могут использоваться только целочисленные значения.

Чаще всего в библиотеках разных языков программирования используются универсальные кучи — бинарные или d-арные. Во-первых, они позволяют извлекать и добавлять элемент за O(log n), а, во-вторых, благодаря свойствам кучи для их хранения достаточно массива элементов. Навигация между родительскими и дочерними элементами для каждого узла происходит по простым формулам.

В .NET 6 в основе PriorityQueue лежит 4-арная min-куча (то есть куча, где каждый узел может иметь 4 элемента уровнем ниже и содержит минимальный элемент своего поддерева), представленная в виде массива кортежей (TElement Element, TPriority Priority)[] _nodes, которая инициализируются пустым массивом по-умолчанию и расширяется при недостаточном Capacity по стандартному правилу List<T> — в 2 раза, но не менее, чем на 4 элемента.

4-арная min-куча всегда удовлетворяет двум свойствам:

1. Свойство формы: 4-арная куча — это полное 4-арное дерево, то есть такое дерево, где все уровни полностью заполнены, за исключением, возможно, последнего уровня. Узлы заполняются слева направо.

2. Свойство кучи: значение (приоритета), хранящееся в каждом узле, меньше или равно его дочерним элементам.

Сложность операций в PriorityQueue связана с сложностью операций в куче:

• Enqueue(TElement, TPriority): O(log₄ N)

• Dequeue(): O(log₄ N)

• Peek(): O(1)

• Count: O(1)

Читать словесное описание алгоритмов кучи — сомнительное удовольствие, а иллюстрации для 4-арной кучи не прибавляют понятности, поэтому иллюстрации к алгоритмам будут на основе биарной min-кучи, а в примерах кода можно увидеть, что работа бинарной и 4-арной кучи отличается только методами навигации по куче, то есть получения индекса родительского и дочерних элементов.

Методы для навигации по куче в PriorityQueue:

private const int Arity = 4; 
private const int Log2Arity = 2;
private static int GetParentIndex(int index) => (index - 1) >> Log2Arity;
private static int GetFirstChildIndex(int index) => (index << Log2Arity) + 1;

Чтобы добавить элемент в кучу, нужно выполнить операцию увеличения кучи (есть разные термины: up-heap, bubble-up, heapify-up):

1. Если куча пуста, то новый элемент становится её корнем (root)

   

2. Добавляем новый элемент на нижний уровень кучи в первое доступное место.

   

3. Сравниваем добавленный элемент с его родителем. Если они в правильном порядке, то операция завершена.

4. Если нет, то меняем добавленный элемент с родительским элементом местами и возвращаемся к предыдущему шагу.

   

private void MoveUpCustomComparer((TElement Element, TPriority Priority) node, int nodeIndex)
{
    while (nodeIndex > 0)
    {
        int parentIndex = GetParentIndex(nodeIndex);
        (TElement Element, TPriority Priority) parent = _nodes[parentIndex];
 
        if (_comparer.Compare(node.Priority, parent.Priority) < 0)
        {
            _nodes[nodeIndex] = parent;
            nodeIndex = parentIndex;
        }
        else
            break;
    }
 
    _nodes[nodeIndex] = node;
}

Количество требуемых операций зависит только от количества уровней, на которые должен подняться новый элемент чтобы привести кучу в согласованное состяоние. 

Итоговая временная сложность в худшем случае составляет O(log N).

Для извлечения минимального значения из min-heap достаточно обратиться к _nodes[0]. Алгоритм удаления корневого (минимального) элемента из кучи:

1. Удаляем значение корневого элемента

   

2. Заменяем текущий удаляемый (корневой) элемент последним элементом кучи (_nodes[0] = _nodes[_size - 1];

   

3. Сравниваем текущий элемент с его потомками. Если они не в правильном порядке (то есть, существуют потомки меньше текущего элемента), то меняем текущий элемент с наименьшим из потомков и повторяем шаг.

   

4. Если они в правильном порядке, то операция завершена.

   

На примере бинарного min-дерева из иллюстрации после извлечения значения [0]-го элемента на его место встает последний ([14]-ый). Из потомков [1]-го и [2]-го выбирается наименьший — [2]-ой, т.к. приоритет [0]-го 26, а [2]-го 18, то они заменяются местами, а [2]-ой элемент становится текущим. Затем из его потомков, [5]-го и [6]-го элементов выбирается тот, у которого приоритет ниже, это [6]-ой элемент. Так как приоритет [2]-го теперь 26, а [6]-го 22, то они меняются местами, а [6]-ой элемент становится текущим. Теперь у текущего элемента только 1 потомок, [13]-ый. Его приоритет выше, чем приоритет текущего, поэтому алгоритм заканчивается.

// node = _nodes[--_size] — последний элемент кучи, 
// nodeIndex = 0 — корневой элемент
private void MoveDownCustomComparer((TElement Element, TPriority Priority) node, int nodeIndex)
{
    int i;
    while ((i = GetFirstChildIndex(nodeIndex)) < size)
    {
        // Ищем дочернюю ноду с минимальным приоритетом
        (TElement Element, TPriority Priority) minChild = _nodes[i];
        int minChildIndex = i;
 
        int childIndexUpperBound = Math.Min(i + Arity, _size);
        while (++i < childIndexUpperBound)
        {
            (TElement Element, TPriority Priority) nextChild = _nodes[i];
            if (_comparer.Compare(nextChild.Priority, minChild.Priority) < 0)
            {
                minChild = nextChild;
                minChildIndex = i;
            }
        }
 
        if (_comparer.Compare(node.Priority, minChild.Priority) <= 0)
        {
            break;
        }
 
        _nodes[nodeIndex] = minChild;
        nodeIndex = minChildIndex;
    }
 
    _nodes[nodeIndex] = node;
}

В PriorityQueue есть пара интересных оптимизаций. Например, при добавлении множества элементов с помощью public void EnqueueRange(IEnumerable<(TElement Element, TPriority Priority)> items!!), фактически коллекция копируется в исходном виде в _nodes, а уже потом массив _nodes один раз перестраивается в кучу с помощью приватного метода Heapify, что позволяет получить итоговую сложность O(log₄ N) вместо O(k * log₄ N), где k — количество добавляемых элементов.

private void Heapify()
{
    int lastParentWithChildren = GetParentIndex(_size - 1);
 
    if (_comparer == null)
        for (int index = lastParentWithChildren; index >= 0; --index)
            MoveDownDefaultComparer(_nodes[index], index);
    else
        for (int index = lastParentWithChildren; index >= 0; --index)
            MoveDownCustomComparer(_nodes[index], index);
}

Метод EnqueueDequeue сравнивает добавляемый элемент с текущим корневым элементов. Если добавляемый элемент имеет меньший приоритет, то он сразу будет извлечен и перестроения кучи не понадобится, а если нет — то вместо двух перестроений для извлечения и добавления элемента понадобится только одно:

public TElement EnqueueDequeue(TElement element, TPriority priority)
{
    // Если очередь пустая, то сразу вернется добавляемый элемент
    if (_size != 0)
    {
        (TElement Element, TPriority Priority) root = _nodes[0];

        if (_comparer == null)
        {
            // Если приоритет добавляемого элемента меньше приоритета корневого,
            // то куча не будет перестраиваться
            if (Comparer<TPriority>.Default.Compare(priority, root.Priority) > 0)
            {
                MoveDownDefaultComparer((element, priority), 0);
                _version++;
                return root.Element;
            }
        }
        else
        {
            // Если приоритет добавляемого элемента меньше приоритета корневого,
            // то куча не будет перестраиваться
            if (_comparer.Compare(priority, root.Priority) > 0)
            {
                MoveDownCustomComparer((element, priority), 0);
                _version++;
                return root.Element;
            }
        }
    }

    return element;
}

Да кто такой этот ваш UnorderedItemsCollection

У новой PriorityQueue есть свойство, которое возвращает элементы очереди в неупорядоченном виде. Но возвращает свойство не массив или енумератор, а вложенный в очередь тип UnorderedItemsCollection. Всё из-за специального енумератора, который отслеживает, что очередь не изменилась во время перебора (не изменилась переменная _version, а перечисление не вышло за пределы очереди) поскольку внутренний массив содержит capacity элементов, а не _size элементов.

public bool MoveNext()
{
    PriorityQueue<TElement, TPriority> localQueue = _queue;
    if (_version == localQueue._version && ((uint)_index < (uint)localQueue._size))
    {
        _current = localQueue._nodes[_index];
        _index++;
        return true;
    }
 
    return MoveNextRare();
}
 
private bool MoveNextRare()
{
    if (_version != _queue._version)
        throw new InvalidOperationException(SR.InvalidOperation_EnumFailedVersion);
 
    _index = _queue._size + 1;
    _current = default;
    return false;
}

Итератор перебирает элементы внутреннего представления кучи _nodes, так что говорить, что UnorderedItemsCollection действительно Unordered не совсем правильно. На самом деле, итератор будет всегда упорядочен в соответствии с правилами хранения 4-арной min-кучи в массиве.

(Очевидные?) ограничения

TPriority может быть любым типом, в том числе изменяемым. Можно, например, сделать такую очередь:

public class TaskPriority
{
    public TaskPriority(int businessValue)
    {
        BusinessValue = businessValue;
    }
 
    public int BusinessValue { get; set; }
}
 
public class TaskPriorityComparer : IComparer<TaskPriority>
{
    public int Compare(TaskPriority x, TaskPriority y) => x.BusinessValue.CompareTo(y.BusinessValue);
}
 
var queue = new PriorityQueue<string, TaskPriority>(new TaskPriorityComparer());

И очередь перестраивает кучу при каждой операции. Но, например, при извлечении элемента куча будет перестраиваться только после извлечения элемента с минимальным приоритетом, даже если приоритет другого элемента изменился и стал минимальным:

var mutablePriority = new TaskPriority(30);
queue.Enqueue("Фронтенд", new TaskPriority(10));
queue.Enqueue("Бэкенд", new TaskPriority(20));
queue.Enqueue("CI", mutablePriority);
 
mutablePriority.BusinessValue = 5;
 
queue.Dequeue(); // (Фронтенд, 10)
queue.Dequeue(); // (CI, 5)
queue.Dequeue(); // (Бэкенд, 20)

При добавлении нового элемента куча перестраивается только до уровня, который нужен добавляемому элементу и всё ломается ещё сильнее:

var mutablePriorityCI = new TaskPriority(30);
var mutablePriorityDocs = new TaskPriority(40);
queue.Enqueue("Фронтенд", new TaskPriority(10));
queue.Enqueue("Бэкенд", new TaskPriority(20));
queue.Enqueue("CI", mutablePriorityCI);
queue.Enqueue("Документация", mutablePriorityCI);
 
mutablePriorityCI.BusinessValue = 5;
mutablePriorityDocs.BusinessValue = 15;
 
queue.Enqueue("Тесты", new TaskPriority(12));
 
queue.Dequeue(); // (Фронтенд, 10)
queue.Dequeue(); // (CI, 5)
queue.Dequeue(); // (Документация, 15)
queue.Dequeue(); // (Тесты, 12)
queue.Dequeue(); // (Бэкенд, 20)

Поэтому ну их изменяемые приоритеты, лучше используйте для TPriority что-нибудь иммутабельное.

Хорошая новость об этом ограничении: в план работ по .NET 7 добавлен issue, позволяющий работать с элементами с изменяемым приоритетом.

Интерактивная иллюстрация бинарной min-кучи

Исходный код PriorityQueue на github