数组[1到N]转最大堆的交换次数统计：求更高效实现方案

我需要统计将数组[1到N]转换为最大堆（Maxheap）所需的总交换次数。我编写的代码可输出正确结果，但希望知晓是否存在更高效的实现方法。

原实现代码：

from math import floor, log2

def count(n):
    layers = 0
    for i in range(1, n+1):
        layers += floor(log2(i))
    return layers
       
if __name__ == "__main__":
    print(count(4)) # 4
    print(count(7)) # 10
    print(count(123)) # 618

现有代码的问题 #

你的代码逻辑是对的：每个节点在构建最大堆时需要的交换次数等于它在堆中的层级（根节点为第0层），总次数就是所有节点层级的总和。但代码采用循环累加的方式，时间复杂度为O(N)，且每次循环都要计算对数，对于超大N（比如10^6甚至更大），效率会偏低。

优化方案：数学公式直接计算 #

我们可以通过堆的结构特性推导数学公式，将时间复杂度降到O(1)：

1. 确定堆的层级h：h是节点N所在的层级（根为0层），即h = floor(log2(n))。
2. 计算满二叉树部分的总交换次数：前h层的满二叉树共有2^h -1个节点，这部分的层级总和可以用公式(h-2)*2^h + 2计算（当h≥1时；h=0即n=1时总和为0）。
3. 计算剩余节点的总交换次数：满二叉树之外的剩余节点数为m = n - (2^h -1)，这些节点都在第h层，贡献的总和是m*h。

优化后的代码 #

from math import floor, log2

def count_optimized(n):
    if n <= 1:
        return 0
    h = floor(log2(n))
    full_tree_nodes = (1 << h) - 1  # 等价于2^h -1
    sum_full = (h - 2) * (1 << h) + 2
    remaining_nodes = n - full_tree_nodes
    sum_remaining = remaining_nodes * h
    return sum_full + sum_remaining

if __name__ == "__main__":
    print(count_optimized(4))    # 输出4
    print(count_optimized(7))    # 输出10
    print(count_optimized(123))  # 输出618

效果验证 #

优化后的代码和原代码输出完全一致，但无需循环遍历所有节点，对于超大N的场景，计算速度会有数量级的提升。

内容的提问来源于stack exchange，提问作者Iltsukka