Python程序员面试算法宝典---解题总结: 第9章 大数据 9.10 如何找出排名前500的数

# -*- coding: utf-8 -*-

'''
Python程序员面试算法宝典---解题总结:  第9章 大数据 9.10 如何找出排名前500的数

题目:
有20个数组，每个数组有500个元素，并且是有序排列好的，现在如何在这20 * 500
个数中找出排名前500的数?

分析:
top K问题，可以设置一个大小为500的小顶堆，
第一次先读取其中1个数组的500个数存入小顶堆，
后续读取其他数组的数字，对每个数字，如果该数字大于堆顶元素，
则删除堆顶元素，添加当前数字，所有数字遍历完，
则堆中的500个元素，就是前500大的元素。
可以用heapq实现，heapq是小顶堆。

关键:
1 书上解法
采用大顶堆加弹出方式，弹出的500个元素即为top 500
假设数组为降序排列，可以设置大顶堆(默认heap为小顶堆，
采用将数字取反放入小顶堆即可)，每次弹出堆顶元素，
然后插入被弹出元素所在数组的下一个元素，调整为小顶堆即可。

2 没有想到
2.1是因为本题求的是大顶堆弹出每次堆顶，直到弹出k个，即为top K。
难点在于heapq默认是小顶堆，所以需要将数字取反，这样将取反后的
数字放入小顶堆，弹出的元素再取反实际就是最大元素。
2.2 没有想到入堆时需要存储的是一个三元组:
(元素，元素所在数组编号，元素在数组中的下标)
这样便于在弹出堆顶元素后，加入被弹出元素所在数组的下一个元素。
2.3 heapq的用法
heap = []
heapq.heappush(heap, arr)
heapq.heappop(heap)
注意堆为空的情况

参考:
Python程序员面试算法宝典
'''

import heapq

def getTopK(data, k):
    if not data or not data[0] or k <= 0:
        return
    arrayNum = len(data)
    arraySize = len(data[0])
    results = []
    # 保持一个最小堆，存放arrayNum个数组中的最小值
    heap = []
    # 初始对每个数组中第一个元素插入到小顶堆中，插入的元素是:(元素，元素所在数组编号，元素在数组中的下标)
    for i in range(0, arrayNum):
        # 注意，这里为了能使用大顶堆，对元素取反插入小顶堆
        value = (-data[i][0], i, 0)
        heapq.heappush(heap, value)
    # 开始尝试弹出
    j = 0
    while j < k:
        # 弹出堆顶元素
        try:
            arr = heapq.heappop(heap)
        except Exception:
            # 说明k超过了所有数组元素之和，这时需要直接退出
            break
        value = -arr[0]
        arrayNum = arr[1]
        index = arr[2]
        results.append(value)
        # 获取被弹出元素所在数组的下一个元素
        if index + 1 < len(data[arrayNum]):
            newArr = (-data[arrayNum][index + 1], arrayNum, index + 1)
            heapq.heappush(heap, newArr)
        else:
            # 说明有一个数组已经全部遍历完成了，则不再追加新的元素
            pass
        j += 1
    return results


def process():
    data = [[29, 17, 14, 2, 1], [19, 17, 16, 15, 6], [30, 25, 20, 14, 5]]
    k = len(data[0])
    results = getTopK(data, k)
    print results
    k = 16
    results = getTopK(data, k)
    print results


if __name__ == "__main__":
    process()