【算法100天 | 17】手撕堆，使插入、删除任意元素的时间复杂度为O(logn)（Java实现）

思路

想要使插入、删除任意元素的时间复杂度为O(logn)，则必须要先用HashMap记住每个元素插入的位置。

数据入堆，就正常查到堆的最后，然后上浮。

数据出堆，需要先在HashMap中找到这个元素在堆中的位置，然后取出堆中的最后一个元素replace和要删除的元素，直接删除堆中要删除的那个元素，heapSize减一；如果这两个数不一样（replace != obj），将replace移动到obj最初的索引位置，然后replace在堆中上浮或下沉，重新维护replace在堆中的索引位置

代码

public class EnhanceHeap<T> {

    private ArrayList<T> heap;
    private HashMap<T, Integer> indexMap;
    private int heapSize;
    private Comparator<? super T> comparator;

    public EnhanceHeap(Comparator<T> comparator) {
        this.comparator = comparator;
        heap = new ArrayList<>();
        heapSize = 0;
        indexMap = new HashMap<>();
    }

    /**
     * 堆顶元素出堆
     */
    public T pop() {
        T heapTop = heap.get(0);
        swap(0, heapSize - 1);
        indexMap.remove(heapTop);
        heap.remove(--heapSize);
        heapify(0);
        return heapTop;
    }

    /**
     * 入堆
     */
    public void push(T obj) {
        heap.add(obj);
        indexMap.put(obj, heapSize);
        heapInsert(heapSize++);
    }

    /**
     * 从堆中删除某一个元素
     */
    public void remove(T obj) {
        Integer index = indexMap.get(obj);
        T replace = heap.get(heapSize - 1);
        // 把要删除的数，和最后一个数都提取出来
        indexMap.remove(obj);
        heap.remove(--heapSize);
        if (obj != replace) {
            heap.set(index, replace);
            indexMap.put(replace, index);
            // 有可能上浮，也有可能下沉
            resign(replace);
        }
    }

    /**
     * obj可能向上、也可能向下调整，但并不知道具体向上还是向下；所以需要调两个，但只有一个会实际产生效果。
     */
    public void resign(T obj) {
        // 向上调整
        heapInsert(indexMap.get(obj));
        // 向下调整
        heapify(indexMap.get(obj));
    }

    /**
     * 数据下沉
     */
    private void heapify(int index) {
        int left = index * 2 + 1;
        while (left < heapSize) {
            // 左子节点小于右子节点吗
            int larger = left + 1 < heapSize && comparator.compare(heap.get(left + 1), heap.get(left)) < 0 ? left + 1 : left;
            larger = comparator.compare(heap.get(larger), heap.get(index)) < 0 ? larger : index;

            if (larger == index) {
                break;
            }
            swap(index, larger);
            index = larger;
            left = index * 2 + 1;
        }
    }

    /**
     * 往堆中插入数据，上浮
     */
    private void heapInsert(int index) {
        // 父节点在前，子节点在后
        while (comparator.compare(heap.get(index), heap.get((index - 1) / 2)) < 0) {
            swap(index, (index - 1) / 2);
            index = (index - 1) / 2;
        }
    }

    private void swap(int i, int j) {
        T o1 = heap.get(i);
        T o2 = heap.get(j);
        heap.set(i, o2);
        heap.set(j, o1);
        indexMap.put(o2, i);
        indexMap.put(o1, j);
    }

    /**
     * 返回堆上的所有元素
     */
    public List<T> getAllElements() {
        List<T> ans = new ArrayList<>();
        for (T c : heap) {
            ans.add(c);
        }
        return ans;
    }

    public boolean isEmpty() {
        return heapSize == 0;
    }

    public int size() {
        return heapSize;
    }

    public boolean contains(T obj) {
        return indexMap.containsKey(obj);
    }

    public T peek() {
        return heap.get(0);
    }

}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

针对基础类型数据的问题

如果堆中的数据是基础类型，同样的数字入堆会有问题；但是可以通过包装一层解决，比如Inner。

class Inner<T> {
    private T value;

    public Inner(T value) {
        this.value = value;
    }
}
1
2
3
4
5
6
7