如何高效计算基于PyWavelets的一维信号SURE阈值？

你当前的代码核心问题在于循环内重复计算平方和，时间复杂度为O(n²)，当n达到50万时会导致严重的性能瓶颈。可以通过以下两种方式大幅提升效率：

方法一：预计算后缀平方和 #

预先计算排序后系数平方的后缀和数组，这样每个i对应的平方和可以直接通过数组索引O(1)获取，将时间复杂度降为O(n)。

优化后的代码：

import numpy as np

# 假设 coeffs、sigma、n 已定义
sorted_coeffs = np.sort(np.abs(np.concatenate(coeffs)))
coeff_sq = sorted_coeffs ** 2

# 计算后缀平方和：suffix_sum[i] 表示从i到末尾的平方和
suffix_sum = np.cumsum(coeff_sq[::-1])[::-1]
# 处理边界：当i等于n时，suffix_sum[n]应为0，所以补一个0
suffix_sum = np.append(suffix_sum, 0)

min_risk = np.inf
T_s = 0.0

for i in range(len(sorted_coeffs)):
    res = (n - i) * sigma**2 + i * (sigma**2 + sorted_coeffs[i]**2)
    risk = suffix_sum[i] - res
    if risk < min_risk:
        min_risk = risk
        T_s = sorted_coeffs[i]

方法二：完全向量化计算（无循环） #

利用numpy的广播特性，一次性计算所有i对应的风险值，再直接取最小值对应的阈值，进一步消除Python循环的开销，时间复杂度仍为O(n)，但实际运行速度会更快。

优化后的代码：

import numpy as np

# 假设 coeffs、sigma、n 已定义
sorted_coeffs = np.sort(np.abs(np.concatenate(coeffs)))
coeff_sq = sorted_coeffs ** 2

# 计算后缀平方和
suffix_sum = np.cumsum(coeff_sq[::-1])[::-1]
suffix_sum = np.append(suffix_sum, 0)

# 生成所有i的数组（0到n-1）
i_arr = np.arange(len(sorted_coeffs))

# 向量化计算所有i对应的res和risk
res = (n - i_arr) * sigma**2 + i_arr * (sigma**2 + coeff_sq)
risk = suffix_sum[:-1] - res  # suffix_sum[:-1]对应i从0到n-1的后缀和

# 找到最小风险对应的索引
min_idx = np.argmin(risk)
min_risk = risk[min_idx]
T_s = sorted_coeffs[min_idx]

优化原理说明 #

• 原代码中每次循环计算sorted_coeffs[j]**2的求和，相当于重复遍历数组，总操作次数是n(n+1)/2，50万数据下就是约1.25e11次操作；
• 预计算后缀平方和只需要一次O(n)的遍历，后续每个i的平方和直接读取数组值；
• 向量化计算则将Python层面的循环转移到numpy的C底层实现，避免了Python循环的额外开销，速度提升更明显。

内容的提问来源于stack exchange，提问作者Владислав Черкасов