""" 在海量数据中选取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])