这个 for 循环可以矢量化吗？

Question

data = np.arange(10)
n = len(data)    
np.asarray([np.sum((data[0:i]-np.mean(data[0:i]))**2) for i in range(1,n)])

这个 for 循环可以通过扩展维度然后折叠它来矢量化吗？

我从某处得到了可以替换的提示

np.mean(data[0:i])

与

np.cumsum(data[0:n-1])/(np.arange(n-1)+1)

Answer 1

它可以按照您的建议通过扩展维度进行矢量化。我认为秘诀是在求和之前使用 np.tril 将级数中的项归零：

# calculate means using cumsum
mean = np.cumsum(data) / np.arange(1, n+1)

# expand into 2 dimensions
mean_2d = np.repeat(mean, n).reshape(n, n)
data_2d = np.tile(data, n).reshape(n, n)

# zero out unneeded terms
diff_squared = np.tril((data_2d-mean_2d)**2)

# sum along rows
np.sum(diff_squared, axis=1)