【Heap实现-java】大根堆/小根堆公用实现代码

大小根堆

• 在二叉树的基础上，大小根堆需要保证层之间的大小关系。

基于数组实现

创建

• 借鉴JDK的PriorityQueue，大小根区分采用Comparator

/**
 * 非线程安全简单版
 *
 * @author SurfingMaster
 * @date 2023/8/20
 */
public class Heap<E> {

    private E arr[];

    private int size;

    private final Comparator<? super E> cmp;

    @SuppressWarnings("all")
    public Heap(int capacity, Comparator<? super E> cmp) {
        this.arr = (E[]) new Object[capacity];
        this.cmp = cmp;
    }

    @SuppressWarnings("all")
    public Heap(E[] initArr, Comparator<? super E> cmp) {
        //this.arr = initArr; // 安全隐患太大，直接指向 initArr 容易与外部引用 冲突

        this.arr = (E[]) new Object[initArr.length];
        System.arraycopy(initArr, 0, this.arr, 0, initArr.length);

        this.cmp = cmp;
        this.size = initArr.length;

        heapify();
    }

    @SuppressWarnings("all")
    public Heap(Collection<E> initColls, Comparator<? super E> cmp) {
        this.arr = (E[]) new Object[initColls.size()];
        this.cmp = cmp;
        size = 0;
        addAll(initColls);
    }

    public void addAll(Collection<E> elements) {
        if (size + elements.size() < arr.length) {
            extendArr(elements.size());
        }
        for (E e : elements) {
            up(e);
            size++;
        }
    }

    /**
     * Floyd
     * 最后一个 非叶子节点
     * int lastIdx = getParentIdx(size - 1);  // 传参调用结果  等价于 (size>>1)-1
     * 由下至上，由右至左 逐个处理
     */
    private void heapify() {
        //for (int i = (size >> 1) - 1; i >= 0; i--) {
        for (int i = getParentIdx(size - 1); i >= 0; i--) {
            down(i);
        }
    }

    public boolean offer(E val) {
        if (isFull()) {
            extendArr();
        }
        // 未自增前的
        //arr[size] = val;
        up(val);
        size++;
        return true;
    }

    @SuppressWarnings("all")
    private final void extendArr() {
        int newLen = (arr.length << 1) - arr.length;    // size
        E[] newArr = (E[]) new Object[newLen];
        System.arraycopy(this.arr, 0, newArr, 0, arr.length);
        this.arr = newArr;
    }

    @SuppressWarnings("all")
    private final void extendArr(int needsSize) {
        int newLen = (arr.length >> 2) + arr.length + needsSize;    // size
        E[] newArr = (E[]) new Object[newLen];
        System.arraycopy(this.arr, 0, newArr, 0, arr.length);
        this.arr = newArr;
    }

    public E peek() {
        if (isEmpty()) return null;
        return arr[0];
    }

    public E poll() {
        if (isEmpty()) return null;
        E res = arr[0];
        //swap(0, size - 1);

        arr[0] = arr[size - 1];
        arr[--size] = null;

        down(0);
        return res;
    }

    private void down(int curIdx) {
        int leftIdx = getLeftIdx(curIdx);
        int rightIdx = leftIdx + 1;

        int stepIdx = curIdx;

        if (leftIdx < size && cmp.compare(arr[leftIdx], arr[stepIdx]) < 0) {
            stepIdx = leftIdx;
        }

        if (rightIdx < size && cmp.compare(arr[rightIdx], arr[stepIdx]) < 0) {
            stepIdx = rightIdx;
        }

        if (stepIdx != curIdx) {
            swap(curIdx, stepIdx);
            down(stepIdx);
        }

    }

    private void up(E val) {
        int cur = size;
        int p = -1;
        while (cur > 0) { // 最顶 就是 根节点
            p = getParentIdx(cur);
            if (cmp.compare(val, arr[p]) < 0) {
                arr[cur] = arr[p];
                //cur = p;
            } else { // todo 简化
                break;
            }
            cur = p;
        }
        arr[cur] = val;
    }

    private int getParentIdx(int idx) {
        return (idx - 1) >> 1;
    }

    private int getLeftIdx(int idx) {
        return (idx << 1) + 1;
    }

    private int getRightIdx(int idx) {
        return (idx << 1) + 2;
    }

    private void swap(int x1, int x2) {
        E t = arr[x1];
        arr[x1] = arr[x2];
        arr[x2] = t;
    }

    public boolean isEmpty() {
        return this.size == 0;
    }

    public boolean isFull() {
        return this.size >= arr.length;
    }

}

使用示例

    @Test
    public void testBuild() {
        List<Integer> initList = List.of(1, 3, 5, 7, 11, 13, 17, 19, 23, 29, 2, 6, 8, 16, 18, 20, 22, 26);

        Heap<Integer> minHeap = new Heap<Integer>(initList, Integer::compare);
        Heap<Integer> maxHeap = new Heap<Integer>(initList, (x1, x2) -> Integer.compare(x2, x1));

    }