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)