如何优化Python中pyramid_blocks函数的执行速度？

我在实现计算金字塔砖块总数的函数时，用循环求和的方式遇到了性能问题，测试提示执行过慢。现有代码如下：

def pyramid_blocks(n, m, h):
    return sum((n+i)*(m+i) for i in range(h))

还有另一种列表实现的版本，性能同样不佳：

def pyramid_blocks(n, m, h):
    r=[]
    t=[]
    mlp=[]
    for i in range(h):
        r.append(n+i)
        t.append(m+i)
    for i, j in zip(r,t):
        mlp.append(i*j)
    return sum(mlp)

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

循环求和的时间复杂度是O(h)，当h很大时会导致执行缓慢。我们可以通过数学推导将计算转为O(1)的公式：

展开每一层的砖块数 (n+i)(m+i) 得到：
nm + (n+m)i + i²

对i从0到h-1求和，利用求和公式：

• 常数项求和：sum(nm) = h * n * m
• 一次项求和：sum((n+m)i) = (n+m) * h*(h-1)/2
• 二次项求和：sum(i²) = h*(h-1)*(2h-1)/6

将三者相加得到总砖块数的公式，用整数运算避免浮点数误差：

def pyramid_blocks(n, m, h):
    return h * n * m + (n + m) * h * (h - 1) // 2 + h * (h - 1) * (2 * h - 1) // 6

这个版本完全消除了循环，无论h多大都能瞬间得到结果，性能大幅提升。

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