"""
在海量数据中选取topK个数据:
整个操作中,遍历数组需要O(n)的时间复杂度,每次调整小顶堆的时间复杂度是O(logK),加起来就是O(nlogK)的复杂度;
如果K远小于n的话,O(nlogK)其实就接近于O(n),甚至会更快,因此也是十分高效的.
"""

# 构建大顶堆
def heapify(arr, n, i, dim):
    largest = i
    left = 2*i + 1
    right = 2*i + 2
    # 与左节点进行比较
    if left < n and arr[i][dim] < arr[left][dim]:
        largest = left
    # 与左节点比较后再与右节点进行比较
    if right < n and arr[largest][dim] < arr[right][dim]:
        largest = right
    # 通过上面跟左右节点比较后，得出三个元素之间较大的下标
    # 若较大下标不是父节点的下标，说明交换后需要重新调整大顶堆
    if largest != i:
        arr[i], arr[largest] = arr[largest], arr[i]
        # 重新调整大顶堆
        heapify(arr, n, largest, dim)

# 堆排序
def heap_sort(arr, topK=None, dim=1): 
    n = len(arr) 
    # 构造大顶堆，从非叶子节点开始倒序遍历，因此是l//2 -1 就是最后一个非叶子节点
    for i in range(n//2-1, -1, -1):
        heapify(arr, n, i, dim)
    # 若topK不为None,则进行设定
    topK = n if topK > n or topK is None else topK
    res_list = []
    # 上面的循环完成了大顶堆的构造，那么就开始把根节点跟末尾节点交换，然后重新调整大顶堆
    for i in range(n-1, n-1-topK, -1):
        res_list.append(arr[0])
        arr[0] = arr[i]
        # arr[i], arr[0] = arr[0], arr[i]
        heapify(arr, i, 0, dim)
    
    return res_list
    # return arr[n-topK:][::-1]

if __name__ == "__main__":
    arr = [[0,12], [1,11], [2,13], [3,5], [4,6], [5,7]]
    arr = heap_sort(arr, 3)
    print ("排序后")
    for i in range(len(arr)):
        print(arr[i])