如何根据大矩阵的第二列制作直方图，然后高效地计算总和？

Question

矩阵就像 4000 列 x 1e8 行。现在我需要根据第二列将 1e8 行分为 1000 个容器，并对每个容器中的行求和。最后，我会得到 4000 列 x 1000 行。

aa=np.random.rand(10000000, 4000)
l,r=min(aa[:,1]),max(aa[:,1])
bins=np.linspace(l,r,1001)
for bn in range(1000):
      ll,rr=bins[bn],bins[bn+1]
      np.sum(dd[(aa[:,1]>ll) &( aa[:,1]< rr) ],axis=0)' # one of the 1000 lines

我发现上面最后一步的速度很低。

然后我需要堆叠所有 1000 行以生成 1000 x 4000 矩阵。

Answer 1

更新后的问题似乎发生了变化，因为我们无法先验地计算行总和，并且箱“标签”没有作为离散值给出（而是需要根据数据计算）。

import numpy as np
rng = np.random.default_rng(83498345935)

# Generate data:
# 1e3 rows (instead of 1e8), 40 columns (instead of 4000), 10 bins (instead of 1000)
n_rows, n_cols, n_bins = 1000, 40, 10
x = rng.random((n_rows, n_cols))

# Compute bins from second column
bin_col = x[:, 1]
lb, ub = bin_col.min(), bin_col.max()
bin_width = (ub - lb) / n_bins
labels = (bin_col - lb) // bin_width
labels[labels == n_bins] = n_bins-1

# Compute sums
import pandas as pd
df = pd.DataFrame(x, index=labels).groupby(level=0).sum()
df.shape  # (10, 40)

---

回答问题的原始理解。

考虑一个数据数组

x

，其中第二列是分箱标签，其余列包含要求和的值：

import numpy as np
rng = np.random.default_rng()
# 1e3 rows (instead of 1e8), 40 columns (instead of 4000)
x = rng.random((1000, 40))  
x[:, 1] = rng.integers(10, size=1000)  # 10 unique values (instead of 1000)

隔离第二列中的 bin 标签，并计算每行中剩余列的总和。确定箱标签并生成箱的边缘。

labels = x[:, 1]
rowsums = x[:, 0] + np.sum(x[:, 2:], axis=1)
bin_labels = np.arange(np.max(labels) + 1)
bin_edges = np.arange(-0.5, np.max(labels) + 1.5)

使用

scipy.stats.binned_statistic

计算分箱总和。与循环参考实现进行比较。

from scipy import stats
res = stats.binned_statistic(labels, rowsums, statistic='sum', bins=bin_edges)
ref = [rowsums[labels==i].sum() for i in bin_labels]
np.testing.assert_allclose(res.statistic, ref)

或者使用

pandas

，你不需要自己生成bin边缘：

import pandas as pd
res2 = pd.Series(rowsums, labels).groupby(level=0).sum()
np.testing.assert_allclose(res2, ref)

如果这两者都太慢或内存密集，我可以发布纯 NumPy 代码，但这些更简单。