Python:第（前）K大（小）数问题

第（前）K大数问题

指的是在长度为n(n>=k)的乱序数组中找出从大到小顺序的第（前）k个数的问题。

算法思路

1. 假设数组长度为N，首先取前K个数，构建最小堆，
2. 将剩余N-K个元素，依次与堆顶元素进行比较，若大于堆顶元素，则替换，并重新为最小堆。

代码

# 构建最小堆
    def min_heap(self, parent, heap):
        child = 2 * parent + 1
        while child < len(heap):
            if child+1 < len(heap) and heap[child+1] < heap[child]:
                child = child+1
            if heap[parent] <= heap[child]:
                break
            heap[parent],heap[child] = heap[child],heap[parent]
            parent, child = child, 2 * child + 1
        return heap
# 找到前k大数
def find_max_kth(self):
        if self.k > len(self.array):
            return None
        heap = self.array[:self.k]
        for i in range(self.k, -1, -1):
            self.min_heap(i, heap)
        for j in range (self.k, len(self.array)):
            if self.array[j] > heap[0]:
                heap[0] = self.array[j]
                self.min_heap(0,heap)
        return heap[0]

第（前）K小数问题

指的是在长度为n(n>=k)的乱序数组中找出从小到大顺序的第（前）k个数的问题。

算法思路

1. 假设数组长度为N，首先取前K个数，构建最大堆，
2. 将剩余N-K个元素，依次与堆顶元素进行比较，若小于堆顶元素，则替换，并重新为最大堆。

# 构建最大堆
    def max_heap(self, parent, heap):
        child = 2 * parent + 1
        while child < len(heap):
            if child+1 < len(heap) and heap[child+1] > heap[child]:
                child = child+1
            if heap[parent] >= heap[child]:
                break
            heap[parent],heap[child] = heap[child],heap[parent]
            parent, child = child, 2 * child + 1
        return heap
# 找到前k小数        
def find_min_kth(self):
        if self.k > len(self.array):
            return None
        heap = self.array[:self.k]
        for i in range(self.k, -1, -1):
            self.max_heap(i, heap)
        for j in range (self.k, len(self.array)):
            if self.array[j] < heap[0]:
                heap[0] = self.array[j]
                self.max_heap(0,heap)
        return heap[0]

完整代码

class Find_Kth(object):

    def __init__(self, k, array):
        self.k = k
        self.array = array

# 构建最小堆
    def min_heap(self, parent, heap):
        child = 2 * parent + 1
        while child < len(heap):
            if child+1 < len(heap) and heap[child+1] < heap[child]:
                child = child+1
            if heap[parent] <= heap[child]:
                break
            heap[parent],heap[child] = heap[child],heap[parent]
            parent, child = child, 2 * child + 1
        return heap

# 构建最大堆
    def max_heap(self, parent, heap):
        child = 2 * parent + 1
        while child < len(heap):
            if child+1 < len(heap) and heap[child+1] > heap[child]:
                child = child+1
            if heap[parent] >= heap[child]:
                break
            heap[parent],heap[child] = heap[child],heap[parent]
            parent, child = child, 2 * child + 1
        return heap

    def find_max_kth(self):
        if self.k > len(self.array):
            return None
        heap = self.array[:self.k]
        for i in range(self.k, -1, -1):
            self.min_heap(i, heap)
        for j in range (self.k, len(self.array)):
            if self.array[j] > heap[0]:
                heap[0] = self.array[j]
                self.min_heap(0,heap)
        return heap[0]

    def find_min_kth(self):
        if self.k > len(self.array):
            return None
        heap = self.array[:self.k]
        for i in range(self.k, -1, -1):
            self.max_heap(i, heap)
        for j in range (self.k, len(self.array)):
            if self.array[j] < heap[0]:
                heap[0] = self.array[j]
                self.max_heap(0,heap)
        return heap[0]

def main():
     orignal_list =[12,4,6,33,54,9,3]
     k = 4
     obj = Find_Kth(k, orignal_list)
     print("第k大数为：", obj.find_max_kth())
     print("第k小数为：", obj.find_min_kth())


 if __name__ == "__main__":
     main()