索引堆（JAVA）

import java.util.Comparator;
import java.util.LinkedList;
import java.util.Objects;
import java.util.Queue;

public class IndexHeap<E> {

    /**
     * demo
     * 某最大堆
     * ********* Z
     * ******* /  \
     * ***** Y     W
     * *** /  \   /  \
     * ** P   O  Q   S
     * * / \
     * A   B
     * 假如这9个元素在数据池中的存储是
     * 数组索引 = [0, 1, 2, 3, 4, 5, 6, 7, 8]
     * data    = [O, W, Z, A, P, S, B, Y, Q]
     * 如果直接将data堆化(heapify) 得到   heap = [Z, Y, W, P, O, Q, S, A, B]
     * 而这里如果只用索引存储堆化后的结果 indexes = [2, 7, 1, 4, 0, 8, 5, 3, 6]
     * 数组索引 = [0, 1, 2, 3, 4, 5, 6, 7, 8]
     * indexes = [2, 7, 1, 4, 0, 8, 5, 3, 6]
     * 之前拿到堆顶元素是 heap[0] = Z, 现在拿到堆顶元素是 data[indexes[0]] = data[2] = 100
     * 已经data中某元素的索引，求堆化后该元素在堆中的索引，rev存储该数据
     * 比如: data中的元素S，其索引为5，堆化后存储在indexes中索引为6的位置，那么 rev[5] = 6
     * data中的元素Q，其索引为8，堆化后存储在indexes中索引为5的位置，那么 rev[8] = 5
     * 数组索引 = [0, 1, 2, 3, 4, 5, 6, 7, 8]
     * rev     = [4, 2, 0, 7, 3, 6, 8, 3, 5]
     */

    // 数据池
    private E[] data;
    // 存储索引的堆
    private int[] indexes;
    // 相当于一个映射 i -> value 表示 data中某元素的index -> 该元素在indexes中的index
    private int[] rev;
    private Comparator<E> comparator;
    private int size;
    private int capacity;

    public IndexHeap(int capacity, Comparator<E> comparator) {
        assert capacity > 0;
        this.capacity = capacity;
        this.size = 0;
        data = (E[]) new Object[capacity];
        indexes = new int[capacity];
        rev = new int[capacity];
        for (int i = 0; i < capacity; i++) {
            indexes[i] = i;
            rev[i] = -1;
        }
        this.comparator = comparator;
    }

    public int size() {
        return this.size;
    }

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

    public void insert(int index, E e) {
        assert this.size < this.capacity;
        assert index >= 0 && index < this.capacity;
        assert Objects.nonNull(e);
        // assert Objects.isNull(data[index]);

        // 赋值
        data[index] = e;

        // 维护索引堆
        indexes[size] = index;
        rev[index] = size;
        shiftUp(size);

        // 维护数量
        this.size++;
    }

    public E peekTop() {
        assert !isEmpty();
        return data[indexes[0]];
    }

    public E extractTop() {
        assert !isEmpty();

        E top = data[indexes[0]];
        swap(indexes, 0, this.size - 1);
        this.size--;
        shiftDown(0);

        return top;
    }

    public int peekTopIndex() {
        assert !isEmpty();
        return indexes[0];
    }

    public int extractTopIndex() {
        assert !isEmpty();

        int topIndex = indexes[0];
        swap(indexes, 0, this.size - 1);
        this.size--;
        shiftDown(0);

        return topIndex;
    }

    public E getElement(int index) {
        assert index >= 0 && index < this.capacity;
        return data[index];
    }

    public void change(int index, E e) {
        assert index >= 0 && index < this.capacity;
        assert Objects.nonNull(e);

        data[index] = e;
        int heapIndex = rev[index];
        shiftUp(heapIndex);
        shiftDown(heapIndex);
    }

    private int parent(int son) {
        return (son - 1) / 2;
    }

    private int leftSon(int parent) {
        return 2 * parent + 1;
    }

    private void shiftUp(int sonIndex) {
        while (sonIndex > 0 && parent(sonIndex) >= 0) {
            int parentIndex = parent(sonIndex);
            E parent = data[indexes[parentIndex]];
            E son = data[indexes[sonIndex]];
            if (comparator.compare(parent, son) >= 0) {
                break;
            }
            swap(indexes, parentIndex, sonIndex);
            sonIndex = parentIndex;
        }
    }

    private void shiftDown(int parentIndex) {
        while (leftSon(parentIndex) < this.capacity) {
            int sonIndex = leftSon(parentIndex);
            if (sonIndex + 1 < this.capacity
                    && comparator.compare(data[indexes[sonIndex + 1]], data[indexes[sonIndex]]) > 0) {
                sonIndex += 1;
            }
            E parent = data[indexes[parentIndex]];
            E son = data[indexes[sonIndex]];
            if (comparator.compare(parent, son) >= 0) {
                break;
            }
            swap(indexes, parentIndex, sonIndex);
            parentIndex = sonIndex;
        }
    }

    private void swap(int[] indexes, int a, int b) {
        int temp = indexes[a];
        indexes[a] = indexes[b];
        indexes[b] = temp;
        // 原本 indexes中: a -> indexes[a]; b -> indexes[b]
        // 原本 rev 中:    indexes[a] -> a; indexes[b] -> b
        // 现在 indexes中: a -> indexes[b]; b -> indexes[a]
        // 现在 rev 中:    indexes[b] -> a; indexes[a] -> b
        // 注意：上边是按原值来分析的
        rev[indexes[a]] = a;
        rev[indexes[b]] = b;
    }

    @Override
    public String toString() {
        if (isEmpty()) {
            return "";
        }
        StringBuilder sb = new StringBuilder();
        Queue<Integer> queue = new LinkedList<>();
        queue.add(0);
        while (!queue.isEmpty()) {
            Integer poll = queue.poll();
            sb.append(data[indexes[poll]]).append(" ");
            if (leftSon(poll) < this.size) {
                queue.add(leftSon(poll));
            }
            if (leftSon(poll) + 1 < this.size) {
                queue.add(leftSon(poll) + 1);
            }
        }
        return sb.toString();
    }
}