python设置堆大小_python数据结构_大顶堆和小顶堆

大顶堆和小顶堆

相关介绍可参看：北京大学空地学院数据结构与算法 第六章 6.8.2.2 小节

代码实现如下

class Heap:

"""二叉堆的实现 小顶堆"""

def __init__(self):

self.heapList = [0] # 默认一个 0 做占位，使得根节点的索引在 1 上

self.currentSize = 0 # 最大节点的索引位置

def perUp(self, i):

"""将小节点逐步上升"""

while i // 2 > 0:

if self.heapList[i] < self.heapList[i // 2]:

self.heapList[i], self.heapList[i // 2] = self.heapList[i // 2], self.heapList[i]

i = i // 2

def insert(self, k):

"""插入节点"""

self.heapList.append(k)

self.currentSize += 1

self.perUp(self.currentSize)

def minChild(self, i):

"""获取左右两个子节点里较小的那个子节点的索引"""

if i * 2 + 1 > self.currentSize: # 右子节点超出节点数量

return i * 2

else:

if self.heapList[i * 2] < self.heapList[i * 2 + 1]:

return i * 2

else:

return i * 2 + 1

def perDown(self, i):

"""将节点下沉到合适位置"""

while (i * 2) <= self.currentSize: # 说明有子节点

mc = self.minChild(i)

if self.heapList[i] > self.heapList[mc]:

self.heapList[i], self.heapList[mc] = self.heapList[mc], self.heapList[i]

i = mc

def delMin(self):

"""删除小节点"""

retval = self.heapList[1] # 删除索引位置为 1 的节点

self.heapList[1] = self.heapList[self.currentSize]

self.heapList.pop()

self.currentSize -= 1

self.perDown(1)

return retval

def buildHeap(self, alist):

i = len(alist) // 2

self.currentSize = len(alist)

self.heapList += alist[:]

while i > 0:

self.perDown(i)

i -= 1

class HeapList(object):

"""大顶推"""

def __init__(self):

self.heaplist = [0]

self.size = 0

def buildHeap(self, alist):

i = len(alist) // 2

self.size = len(alist)

self.heaplist += alist[:]

while i > 0:

self.percDown(i)

i -= 1

def percUp(self, i):

while i // 2 > 0:

if self.heaplist[i] > self.heaplist[i // 2]:

self.heaplist[i], self.heaplist[i // 2] = self.heaplist[i // 2], self.heaplist[i]

i //= 2

def insert(self, k):

self.heaplist.append(k)

self.size += 1

self.percUp(self.size)

def maxChild(self, i):

if i * 2 + 1 > self.size:

return i * 2

else:

if self.heaplist[i * 2] > self.heaplist[i * 2 + 1]:

return i * 2

else:

return i * 2 + 1

def percDown(self, i):

while i * 2 <= self.size:

mc = self.maxChild(i)

if self.heaplist[i] < self.heaplist[mc]:

self.heaplist[i], self.heaplist[mc] = self.heaplist[mc], self.heaplist[i]

i = mc

def delMax(self):

retval = self.heaplist[1]

self.heaplist[1] = self.heaplist[self.size]

self.size -= 1

self.heaplist.pop()

self.percDown(1)

return retval

# 采用大顶堆的方式，制作容量为 k 的大顶堆，向堆中添加元素时，比堆顶值小，就弹出堆顶，并将此元素添加进堆。这就保证，最后遍历完成后，

# 我们获得了比堆顶小的 k-1 个最小值

# 时间复杂度 O(nlogK) 因为只维护 K 大小的堆

class Solution:

def getLeastNumbers(self, arr, k):

if k == 0:

return []

heaplist = HeapList()

heaplist.buildHeap(arr[:k])

for i in arr[k: ]:

if i < heaplist.heaplist[1]:

heaplist.delMax()

heaplist.insert(i)

return heaplist.heaplist[1:]

if __name__ == '__main__':

solution = Solution()

arrlist = [1, 2, 3, 4, 5, 6, 7, 8]

res = solution.getLeastNumbers(arrlist, 3)

print(res)