如何用无循环Numpy实现一维数组与递增序列的逐步滑动相关计算

问题描述 #

给定一维数组，需要依次计算前2个、前3个……直至全部元素与对应长度的递增序列（如np.arange(k)，k为当前段长度）的皮尔逊相关系数。原实现使用循环结合Numpy，现要求用更高效的无循环Numpy方法实现。

原循环代码：

import numpy as np

data = np.array([10,5,8,9,15,22,26,11,15,16,18,7,4,8,-2,-3,-4,-6,-2,0,10,0,5,8])
correl = np.zeros(data.shape)
for i in range(1, data.shape[0]):
     correl[i] = np.corrcoef(data[0: i+1], np.arange(i+1))[0, 1]
print(correl)

原输出：

[ 0.    -1.    -0.397  0.     0.607  0.799  0.88   0.64   0.581  0.556
  0.574  0.322  0.078 -0.02  -0.237 -0.383 -0.489 -0.572 -0.614 -0.634
 -0.568 -0.59  -0.573 -0.533]

高效无循环实现 #

利用皮尔逊相关系数的代数表达式，结合Numpy的累积求和函数np.cumsum实现全量向量化计算，避免循环开销：

import numpy as np

data = np.array([10,5,8,9,15,22,26,11,15,16,18,7,4,8,-2,-3,-4,-6,-2,0,10,0,5,8])
n = len(data)
y = np.arange(n)  # 对应各段的递增序列

# 计算各项累积和
sum_x = np.cumsum(data)
sum_y = np.cumsum(y)
sum_xy = np.cumsum(data * y)
sum_x2 = np.cumsum(data ** 2)
sum_y2 = np.cumsum(y ** 2)
k = np.arange(1, n+1)  # 各段的长度

# 计算相关系数的分子和分母
numerator = k * sum_xy - sum_x * sum_y
denominator_x = k * sum_x2 - sum_x ** 2
denominator_y = k * sum_y2 - sum_y ** 2
denominator = np.sqrt(denominator_x * denominator_y)

# 计算相关系数，处理k=1时分母为0的情况
correl = numerator / denominator
correl[0] = 0.0  # 对应原代码中第一个元素为0

# 保留三位小数以匹配原输出格式
print(np.round(correl, 3))

运行输出：

[ 0.    -1.    -0.397  0.     0.607  0.799  0.88   0.64   0.581  0.556
  0.574  0.322  0.078 -0.02  -0.237 -0.383 -0.489 -0.572 -0.614 -0.634
 -0.568 -0.59  -0.573 -0.533]

原理说明 #

皮尔逊相关系数可通过以下公式计算：
[
r = \frac{k\sum xy - \sum x \sum y}{\sqrt{(k\sum x^2 - (\sum x)^2)(k\sum y^2 - (\sum y)^2)}}
]
其中k为当前计算段的元素数量。

通过np.cumsum一次性计算所有前缀的累积和，再利用Numpy的广播特性完成向量化运算，相比循环实现，时间复杂度从O(n²)降至O(n)，效率大幅提升，尤其适合处理大规模数组。

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