python 堆排序模板

# 调整某一个节点到具体的位置，确保堆顶为最大值,时间复杂度为logn
def heapify(arr, idx, end):
    # size = len(arr)
    left = idx * 2 + 1
    right = left + 1
    while left <= end:  # 当前节点非叶子结点
        max_idx = idx
        if arr[left] > arr[max_idx]:
            max_idx = left
        if right <= end and arr[right] > arr[max_idx]:
            max_idx = right
        if idx == max_idx:  # 如果不用交换则break，说明已经交换结束
            break
        arr[idx], arr[max_idx] = arr[max_idx], arr[idx]  # 将父子三个节点中，最大的值放到父节点
        idx = max_idx  # 继续调整子树
        left = idx * 2 + 1
        right = left + 1


# 让一个无序数组变成大根堆， 时间复杂度为n*logn
def build_heap(arr):
    size = len(arr)
    for i in range((size - 2) // 2, -1, -1):  # （size-2）//2 是最后一个非叶节点，叶节点不用调整
        heapify(arr, i, size - 1)
    return arr


# 升序堆排序，思路为先建堆，然后让堆顶最大元素和最后一个交换，然后继续调整第一个到倒数第二个元素，获取最大值，再交换
# 时间复杂度为n*logn + n*logn = n*logn
def heap_sort(arr):
    build_heap(arr)
    size = len(arr)
    for i in range(size):
        arr[0], arr[size - i - 1] = arr[size - i - 1], arr[0]
        heapify(arr, 0, size - i - 2)
    return arr


print(heap_sort([3, 2, 1]))