Skip to main content

Module 0x2::priority_queue

Priority queue implemented using a max heap.

use 0x1::vector;

Struct PriorityQueue​

Struct representing a priority queue. The entries vector represents a max heap structure, where entries[0] is the root, entries[1] and entries[2] are the left child and right child of the root, etc. More generally, the children of entries[i] are at i * 2 + 1 and i * 2 + 2. The max heap should have the invariant that the parent node's priority is always higher than its child nodes' priorities.

struct PriorityQueue<T: drop> has drop, store
Fields

Struct Entry​

struct Entry<T: drop> has drop, store
Fields
priority: u64
value: T

Constants​

For when heap is empty and there's no data to pop.

const EPopFromEmptyHeap: u64 = 0;

Function new​

Create a new priority queue from the input entry vectors.

public fun new<T: drop>(entries: vector<priority_queue::Entry<T>>): priority_queue::PriorityQueue<T>
Implementation
public fun new<T: drop>(mut entries: vector<Entry<T>>): PriorityQueue<T> {
    let len = entries.length();
    let mut i = len / 2;
    // Max heapify from the first node that is a parent (node at len / 2).
    while (i > 0) {
        i = i - 1;
        max_heapify_recursive(&mut entries, len, i);
    };
    PriorityQueue { entries }
}

Function pop_max​

Pop the entry with the highest priority value.

public fun pop_max<T: drop>(pq: &mut priority_queue::PriorityQueue<T>): (u64, T)
Implementation
public fun pop_max<T: drop>(pq: &mut PriorityQueue<T>): (u64, T) {
    let len = pq.entries.length();
    assert!(len > 0, EPopFromEmptyHeap);
    // Swap the max element with the last element in the entries and remove the max element.
    let Entry { priority, value } = pq.entries.swap_remove(0);
    // Now the max heap property has been violated at the root node, but nowhere else
    // so we call max heapify on the root node.
    max_heapify_recursive(&mut pq.entries, len - 1, 0);
    (priority, value)
}

Function insert​

Insert a new entry into the queue.

public fun insert<T: drop>(pq: &mut priority_queue::PriorityQueue<T>, priority: u64, value: T)
Implementation
public fun insert<T: drop>(pq: &mut PriorityQueue<T>, priority: u64, value: T) {
    pq.entries.push_back(Entry { priority, value });
    let index = pq.entries.length() - 1;
    restore_heap_recursive(&mut pq.entries, index);
}

Function new_entry​

public fun new_entry<T: drop>(priority: u64, value: T): priority_queue::Entry<T>
Implementation
public fun new_entry<T: drop>(priority: u64, value: T): Entry<T> {
    Entry { priority, value }
}

Function create_entries​

public fun create_entries<T: drop>(p: vector<u64>, v: vector<T>): vector<priority_queue::Entry<T>>
Implementation
public fun create_entries<T: drop>(mut p: vector<u64>, mut v: vector<T>): vector<Entry<T>> {
    let len = p.length();
    assert!(v.length() == len, 0);
    let mut res = vector[];
    let mut i = 0;
    while (i < len) {
        let priority = p.remove(0);
        let value = v.remove(0);
        res.push_back(Entry { priority, value });
        i = i + 1;
    };
    res
}

Function restore_heap_recursive​

fun restore_heap_recursive<T: drop>(v: &mut vector<priority_queue::Entry<T>>, i: u64)
Implementation
fun restore_heap_recursive<T: drop>(v: &mut vector<Entry<T>>, i: u64) {
    if (i == 0) {
        return
    };
    let parent = (i - 1) / 2;

    // If new elem is greater than its parent, swap them and recursively
    // do the restoration upwards.
    if (*&v[i].priority > *&v[parent].priority) {
        v.swap(i, parent);
        restore_heap_recursive(v, parent);
    }
}

Function max_heapify_recursive​

Max heapify the subtree whose root is at index i. That means after this function finishes, the subtree should have the property that the parent node has higher priority than both child nodes. This function assumes that all the other nodes in the subtree (nodes other than the root) do satisfy the max heap property.

fun max_heapify_recursive<T: drop>(v: &mut vector<priority_queue::Entry<T>>, len: u64, i: u64)
Implementation
fun max_heapify_recursive<T: drop>(v: &mut vector<Entry<T>>, len: u64, i: u64) {
    if (len == 0) {
        return
    };
    assert!(i < len, 1);
    let left = i * 2 + 1;
    let right = left + 1;
    let mut max = i;
    // Find the node with highest priority among node `i` and its two children.
    if (left < len && *&v[left].priority > *&v[max].priority) {
        max = left;
    };
    if (right < len && *&v[right].priority > *&v[max].priority) {
        max = right;
    };
    // If the parent node (node `i`) doesn't have the highest priority, we swap the parent with the
    // max priority node.
    if (max != i) {
        v.swap(max, i);
        // After the swap, we have restored the property at node `i` but now the max heap property
        // may be violated at node `max` since this node now has a new value. So we need to now
        // max heapify the subtree rooted at node `max`.
        max_heapify_recursive(v, len, max);
    }
}

Function priorities​

public fun priorities<T: drop>(pq: &priority_queue::PriorityQueue<T>): vector<u64>
Implementation
public fun priorities<T: drop>(pq: &PriorityQueue<T>): vector<u64> {
    let mut res = vector[];
    let mut i = 0;
    while (i < pq.entries.length()) {
        res.push_back(pq.entries[i].priority);
        i = i +1;
    };
    res
}