如何利用预计算分箱索引向量化实现分箱数据均值计算？

问题背景 #

假设我有一个包含多个观测值（a和b）的时间序列（t），数据生成代码如下：

import numpy as np

t = np.linspace(0, 10, 100)
a = np.random.normal(loc=5, scale=0.1, size=t.size)
b = np.random.normal(loc=1, scale=0.5, size=t.size)

现有循环版分箱均值实现 #

想要获取时间分箱的均值，当前使用循环的实现代码如下：

bin_edges = np.linspace(0, 12, 12)
bin_index = np.digitize(t, bin_edges) - 1
a_binned = np.zeros(bin_edges.size - 1)
b_binned = np.zeros(bin_edges.size - 1)

for ibin in np.argwhere(np.bincount(bin_index) > 0).flatten():
    select = bin_index == ibin
    a_binned[ibin] = np.mean(a[select])
    b_binned[ibin] = np.mean(b[select])

问题 #

如何将上述循环逻辑向量化，提升计算效率？

向量化解决方案 #

方法一：利用np.bincount直接计算分箱均值 #

np.bincount可以一次性计算每个分箱的样本总和与样本数量，再通过除法得到均值，完全避免循环：

# 分箱设置（和原代码一致）
bin_edges = np.linspace(0, 12, 12)
bin_index = np.digitize(t, bin_edges) - 1
n_bins = bin_edges.size - 1

# 计算每个分箱的总和
a_sum = np.bincount(bin_index, weights=a, minlength=n_bins)
b_sum = np.bincount(bin_index, weights=b, minlength=n_bins)

# 计算每个分箱的样本数
bin_counts = np.bincount(bin_index, minlength=n_bins)

# 计算均值，同时处理空分箱（避免除以0）
a_binned = np.where(bin_counts > 0, a_sum / bin_counts, 0)
b_binned = np.where(bin_counts > 0, b_sum / bin_counts, 0)

方法二：多变量批量处理 #

如果需要处理多个观测变量（比如a、b甚至更多），可以将数据合并后批量计算，减少重复代码：

# 将a和b合并为二维数组（每行对应一个时间点的观测值）
data = np.column_stack([a, b])

# 批量计算每个变量的分箱总和
sum_binned = np.array([
    np.bincount(bin_index, weights=data[:, col], minlength=n_bins)
    for col in range(data.shape[1])
])

# 批量计算均值
mean_binned = np.where(bin_counts > 0, sum_binned / bin_counts, 0)

# 拆分回单个变量的分箱均值
a_binned, b_binned = mean_binned[0], mean_binned[1]

结果验证 #

可以通过对比循环版结果验证向量化实现的正确性：

# 循环版结果（原代码逻辑）
a_binned_loop = np.zeros(n_bins)
b_binned_loop = np.zeros(n_bins)
for ibin in np.argwhere(np.bincount(bin_index) > 0).flatten():
    select = bin_index == ibin
    a_binned_loop[ibin] = np.mean(a[select])
    b_binned_loop[ibin] = np.mean(b[select])

# 验证数值一致性
print(np.allclose(a_binned, a_binned_loop))  # 输出 True
print(np.allclose(b_binned, b_binned_loop))  # 输出 True

内容的提问来源于stack exchange，提问作者Tom de Geus