温习Algs4 (二):优先级队列(堆)

优先级队列

优先级队列 (Priority Queue) 实质上是一个堆, 支持的操作是push和pop, pop时总是返回当前队列中的最小(大)值.

PQ.java

/******************************************************************************
 *  Compilation:  javac PQ.java
 *  Execution:    java PQ
 *  Author:       Chenghao Wang
 ******************************************************************************/

public class PQ<T extends Comparable<T>> {
    public static final int MIN_PQ = 0;
    public static final int MAX_PQ = 1;
    private Vector<T> heap = new Vector<T>();
    private int type;

    PQ() {
        this(MIN_PQ);
    }

    PQ(int type) {
        this.type = type;
        heap.add(null);
    }

    private boolean shouldSwap(int up, int down) {
        T a = heap.get(up);
        T b = heap.get(down);
        int cmp = a.compareTo(b);
        if (type == MIN_PQ) {
            return cmp > 0;
        } 
        else {
            return cmp < 0;
        }
    }

    private void swap(int i, int j) {
        T tmp = heap.get(i);
        heap.set(i, heap.get(j));
        heap.set(j, tmp);
    }

    public void push(T item) {
        heap.add(item);
        for (int curr = heap.size() - 1; curr > 1; curr /= 2) {
            int next = curr / 2;
            if (shouldSwap(next, curr)) {
                swap(next, curr);
            } 
            else {
                break;
            }
        }
    }

    public T pop() {
        T result = heap.get(1);
        swap(1, heap.size() - 1);
        heap.remove(heap.size() - 1);
        int next = 0;
        for (int curr = 1; (curr * 2) < heap.size(); curr = next) {
            next = curr * 2;
            if (next < (heap.size() - 1) && shouldSwap(next, next + 1)) {
                next++;
            }
            if (shouldSwap(curr, next)) {
                swap(curr, next);
            } 
            else {
                break;
            }
        }
        return result;
    }
}

索引优先级队列

索引优先级队列 (Indexed Priority Queue) 是在优先级队列的基础上添加了索引功能, 即可以将一个值与一个索引绑定, 主要用于 Dijkstra和最小生成树算法上.

IndexPQ.java

/******************************************************************************
 *  Compilation:  javac IndexPQ.java
 *  Execution:    java IndexPQ
 *  Author:       Chenghao Wang
 ******************************************************************************/

public class IndexPQ<T extends Comparable<T>> {
    public static final int MIN_PQ = 0;
    public static final int MAX_PQ = 1;
    private boolean[] mark;
    private int[] heap;
    private int heapPtr;
    private T[] indexes;
    private int type;
    private int capacity;

    IndexPQ(int size) {
        this(MIN_PQ, size);
    }

    IndexPQ(int type, int size) {
        capacity = size;
        this.type = type;
        heap = new int[size + 1];
        mark = new boolean[size];
        indexes = (T[]) (new Comparable[size]);
        heapPtr = 1;
    }

    private boolean shouldSwap(int up, int down) {
        T a = indexes[heap[up]];
        T b = indexes[heap[down]];
        int cmp = a.compareTo(b);
        if (type == MIN_PQ) {
            return cmp > 0;
        } 
        else {
            return cmp < 0;
        }
    }

    private void swap(int i, int j) {
        int tmp = heap[i];
        heap[i] = heap[j];
        heap[j] = tmp;
    }

    private void adjustBottomUp(int i) {
        for (int curr = i; curr > 1; curr /= 2) {
            int next = curr / 2;
            if (shouldSwap(next, curr)) {
                swap(next, curr);
            } 
            else {
                break;
            }
        }
    }

    private void adjustTopDown(int i) {
        int next  = 1;
        for (int curr = i; (curr * 2) < heapPtr; curr = next) {
            next = curr * 2;
            if ((next < (heapPtr - 1)) && shouldSwap(next, next + 1)) {
                next++;
            }
            if (shouldSwap(curr, next)) {
                swap(curr, next);
            } 
            else {
                break;
            }
        }
    }

    public void push(int index, T item) {
        if (mark[index]) {
            change(index, item);
            return;
        }

        mark[index] = true;
        heap[heapPtr++] = index;
        indexes[index] = item;
        adjustBottomUp(heapPtr - 1);
    }

    public void change(int index, T item) {
        T old = indexes[index];
        int cmp = old.compareTo(item);
        if (cmp == 0) return;
        indexes[index] = item;
        int i = 1;
        while (heap[i] != index) i++;
        if (type == MIN_PQ) {
            if (cmp < 0) {
                adjustTopDown(i);
            }
            else {
                adjustBottomUp(i);
            }
        }
        else {
            if (cmp > 0) {
                adjustTopDown(i);
            }
            else {
                adjustBottomUp(i);
            }        
        }
    }

    public int pop() {
        int result = heap[1];
        mark[result] = false;
        heapPtr--;
        swap(1, heapPtr);
        adjustTopDown(1);
        return result;
    }

    public T getValue(int index) {
        return indexes[index];
    }

    public boolean contains(int index) {
        return mark[index];
    }

    public boolean isEmpty() {
        return heapPtr == 1;
    }

    public int size() {
        return heapPtr - 1;
    }
}

总结

这两个数据结构我认为非常有用, 不仅可以代替排序算法, 而且在后续的最短路径, 最小生成树和最小费用最大流问题上扮演着关键的角色.