堆的python实现_堆实现代码python

from binary_tree import Node

# 完全二叉树的叶子节点只存在于最下面两层
# 且最下面一层的叶子节点全部靠左紧密排列
# 节点编号为从0开始，从上往下，从左到右依次编号
# 某个节点i的父节点是（i-1)//2，左节点是2*i+1,有节点是2*i+2
class PerfectBinaryTree(Node):
    def is_empty(self):
        if not self.item:
            return True

    def add_one(self, data):
        if self.is_empty():
            self.item = data
            return
        queue = [self]
        while queue:
            size = len(queue)
            for _ in range(size):
                cur = queue.pop(0)
                if cur.lchild:
                    queue.append(cur.lchild)
                else:
                    cur.lchild = Node(data)
                    return

                if cur.rchild:
                    queue.append(cur.rchild)
                else:
                    cur.rchild = Node(data)
                    return

    def add_batch(self, arr):
        if self.is_empty():
            self.item = arr[0]
            arr.pop(0)
        queue = [self]
        while arr and queue:
            size = len(queue)
            for _ in range(size):
                cur = queue.pop(0)
                if cur.lchild:
                    queue.append(cur.lchild)
                elif arr:
                    cur.lchild = Node(arr[0])
                    arr.pop(0)
                    queue.append(cur.lchild)

                if cur.rchild:
                    queue.append(cur.rchild)
                elif arr:
                    cur.rchild = Node(arr[0])
                    arr.pop(0)
                    queue.append(cur.rchild)

    def bfs(self):
        queue = [self]
        while queue:
            size = len(queue)
            for _ in range(size):
                cur = queue.pop(0)
                if cur.item:
                    print(cur.item)
                if cur.lchild:
                    queue.append(cur.lchild)
                if cur.rchild:
                    queue.append(cur.rchild)



from perfect_binary_tree import PerfectBinaryTree

"""
二叉堆首先是一颗完全二叉树（结构上）
二叉堆节点间关系满足：
堆中根节点<=子树中的节点值（最小堆）
堆中根节点>=子树中的节点值（最大堆）
注意：在最大堆中，只能保证当前根节点大于等于子树的所有节点（任意子树都满足），节点大小与所处层数无关
"""
class Heap(PerfectBinaryTree):
    def __init__(self, type='big'):
        self.type = type
        self.arr = []
        if self.type == 'big':
            self._operator = lambda a, b: a > b
        else:
            self._operator = lambda a, b: a < b

    def visualize(self):
        super().__init__()
        super().add_batch(self.arr.copy())

    def _swap(self, i, j):
        self.arr[i], self.arr[j] = self.arr[j], self.arr[i]

    # //为向下取整
    def _get_parent_index(self, i):
        return (i - 1) // 2

    def _get_lchild_index(self, i):
        return 2 * i + 1

    def _get_rchild_index(self, i):
        return 2 * i + 2

    def _get_right_bottom_most_nonleaf_index(self):
        return self._get_parent_index(len(self.arr) - 1)

    def add_one(self, data):
        """
        上浮操作发生在追加节点的时候，在数据的最后加节点，由于必须保证堆的性质，因此该节点要上浮
        """
        self.arr.append(data)
        self._siftup(len(self.arr) - 1)

    def _siftup(self, i):
        # use this function only when self.arr is already heapified
        # i = 0 indicates root. Sift up until root
        while (i > 0) and (self._operator(self.arr[i], self.arr[self._get_parent_index(i)])):
            self._swap(i, self._get_parent_index(i))
            i = self._get_parent_index(i)

    def add_batch(self, arr):
        """
        下浮发生在删除根节点的时候，此时把最后一个元素挪到根节点上，因此要下沉
        或者根节点被替换的时候，也要下沉

        下浮法建堆的复杂度：
        从最下面最右边的非叶子节点开始遍历到根节点
        设完全二叉树总共有 h 层，如果当前遍历节点位于第 k 层，则需要下浮 h-k 次，这一层有 2^(k-1)个节点
        则这一层总共需要 (h-k) *  2^(k-1) 次下层操作
        遍历需要 S = (h-(h-1)) * 2^(h-1-1) + 2 * 2^(h-3) + 3 * 2^(h-4) .... (h-1) * 2^0
                  = 1 * 2^(h-2) + 2 * 2^(h-3) + .... (h-1) * 2^0
               2S = 2 * 2^(h-1) + 2 * 2^(h-2) + .... (h-1) * 2
        2S-S = S  = 2 * 2^(h-1) + [2^(h-2) + 2^(h-3) + ... 2^1] - (h-1) * 2^0
                  = 2 * 2^(h-1) - (h-1) * 2^0 + [2^(h-2) + 2^(h-3) + ... 2^1]
                  = 2 * 2^(h-1) - (h-1) * 2^0 + O(2^h)
                  = O(2^h)
                  = O(n)
        """
        self.arr += arr
        # heapify the arr
        for i in range(self._get_right_bottom_most_nonleaf_index(), -1, -1):
            self._siftdown(i, len(self.arr)-1)

    def _siftdown(self, i, max_index):
        # use this function only when self.arr is already heapified
        # Sift down until max_index
        while i <= max_index:
            next_index = i
            if (self._get_lchild_index(i) <= max_index) and \
                    (self._operator(self.arr[self._get_lchild_index(i)], self.arr[i])):
                next_index = self._get_lchild_index(i)
            if (self._get_rchild_index(i) <= max_index) and \
                    (self._operator(self.arr[self._get_rchild_index(i)], self.arr[next_index])):
                next_index = self._get_rchild_index(i)

            if next_index != i:
                self._swap(i, next_index)
                i = next_index
            else:
                break

    def sort(self):
        """
        从后往前遍历到第 i个节点时，从i开始到堆顶的堆高为log(i), 因此堆顶需要下降 log(i)次
        总共次数 S = log(n-1) + log(n-2) + ..log(i).. + 1
                 <= integral~[1, n-1](logx)    注意 logx 的原函数为 logx * x - x
                  = logx * x - x |x=(n-1)  -  log x * x - x |x=1
                  = O(log(n-1) * (n-1))
                  = O(logn * n)
        """
        # note that the sorted self.arr is in reverse relationship with self.type
        for max_index in range(len(self.arr) - 1, 0, -1):
            self._swap(0, max_index)
            self._siftdown(0, max_index-1)




arr = [1, 2, 2, 3, 4, 5, 6, 7, 8, 8, 9, 10]
a = Heap(type='small')
a.add_batch(arr[:first_x])
# complexity O(n)
for i in range(first_x, len(arr)):
    if arr[i] > a.arr[0]:
        a.arr[0] = arr[i]
        # complexity: O(log(first_x))
        a._siftdown(0, first_x-1)
# total complexity: O(n * log(first_x))
print(a.arr)

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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192