python 实现堆,优先队列----处理海量数据的topK问题

2020-02-16  本文已影响0人  金融测试民工

    堆 处理海量数据的topK,分位数非常合适,优先队列应用在元素优先级排序。比如数组的频率排序非常合适。与基于比较的排序算法 时间复杂度O(nlogn) 相比, 使用堆,优先队列复杂度可以下降到 O(nlogk),在总体数据规模 n 较大,而维护规模 k 较小时,时间复杂度优化明显。

    堆,优先队列的本质其实就是个完全二叉树,有其下重要性质:

1、父节点index为 (i-1) // 2

2、左子节点index为 2*i + 1

3、右子节点index为 2*i + 2

4、大顶堆中每个父节点大于子节点,小顶堆每个父节点小于子节点

5、优先队列以优先级为堆的排序依据

    因为性质1,2,3,堆可以用数组直接来表示,不需要通过链表建树。

    堆,优先队列 有两个重要操作,时间复杂度均是 O(logk)。以大顶锥为例:

1、上浮sift up: 向堆新加入一个元素,堆规模+1,依次向上与父节点比较,如大于父节点就交换。

2、下沉sift down: 从堆取出一个元素(堆规模-1,用于堆排序)或者更新堆中一个元素(本题),逆序遍历数组index从 (k-1) // 2 到 index为 0,向下走保证父节点大于子节点。

    对于topk 问题:最大堆求topk小,最小堆求topk大。

    topk小:构建一个k个数的最大堆,当读取的数小于根节点时,替换根节点,重新塑造最大堆

    topk大:构建一个k个数的最小堆,当读取的数大于根节点时,替换根节点,重新塑造最小堆

这一题的总体思路 总体时间复杂度 O(nlogk)

~建立字典遍历一次统计出现频率. O(logn)

~取前k个数,构造规模为k的最小堆 minheap. O(logn)

~遍历规模k之外的数据,大于堆顶则入堆,维护规模为k的最小堆 minheap. O(nlogk)

~(如需按频率输出,对规模为k的堆进行排序)

class Solution:

    def topKFrequent(self, nums: List[int], k: int) -> List[int]:

        # hashmap 统计频率

        freq_count = {}

        for num in nums:

            if num in freq_count:

                freq_count[num] += 1

            else:

                freq_count[num] = 1

        def sift_up(arr, k):

            """ 时间复杂度 O(logk) k 为堆的规模"""

            new_index, new_val = k-1, arr[k-1]

            while (new_index > 0 and arr[(new_index-1)//2][1] > new_val[1]):

                arr[new_index] = arr[(new_index-1)//2]

                new_index = (new_index-1)//2

            arr[new_index] = new_val # 这里采用的是类似插入排序的赋值交换

        def sift_down(arr, root, k):

            """ O(logk). 右节点index 2*root+1,左节点 2*root+1, 父节点 (child-1)//2"""

            root_val = arr[root]

            while (2*root+1 < k):

                # 右节点 2*root+1,左节点 2*root+1, 父节点 (child-1)//2

                child = 2 * root + 1

                if child+1 < k and arr[child][1] > arr[child+1][1]:

                    child += 1 # 左右节点中最大的对应index

                # 小顶锥 用 >,大顶锥 用 <

                if root_val[1] > arr[child][1]:

                    arr[root] = arr[child]

                    root = child # 继续向下检查

                else: break # 如果到这里没乱序,不用再检查后续子节点

            arr[root] = root_val

        # 注意构造规模为k的堆, 时间复杂度O(n),因为堆的规模是从0开始增长的

        freq_list = list(freq_count.items())

        min_heap = []

        for i in range(k):

            min_heap.append(freq_list[i])

            sift_up(min_heap, i+1) 

        # 遍历剩下元素,大于堆顶入堆,下沉维护小顶堆

        for item in freq_list[k:]:

            priority = item[1]

            if priority > min_heap[0][1]:

                min_heap[0] = item

                sift_down(min_heap, 0, k)

        return [item[0] for item in min_heap]

再附上堆排序,其实就是

    1、下沉(堆规模不变),构造大顶锥

    2、下沉(堆规模-1),依次把最大值放到尾部,不再维护。

def heapSort(arr):

    def sift_down(arr, root, k):

        root_val = arr[root] # 用插入排序的赋值交换

        # 确保交换后,对后续子节点无影响

        while (2*root+1 < k):

            # 构造根节点与左右子节点

            child = 2 * root + 1  # left = 2 * i + 1, right = 2 * i + 2

            if child+1 < k and arr[child] < arr[child+1]: # 如果右子节点在范围内且大于左节点

                child += 1

            if root_val < arr[child]:

                arr[root] = arr[child]

                root = child

            else: break # 如果有序,后续子节点就不用再检查了

        arr[root] = root_val

    k = len(arr) # k 为heap的规模

    # 构造 maxheap. 从倒数第二层起,该元素下沉

    for i in range((k-1)//2, -1, -1):

        sift_down(arr, i, k)

    # 从尾部起,依次与顶点交换并再构造 maxheap,heap规模-1

    for i in range(k - 1, 0, -1):

        arr[i], arr[0] = arr[0], arr[i]  # 交换

        sift_down(arr, 0, i)

上一篇下一篇

猜你喜欢

热点阅读