首页 > 解决方案 > NumPy - 使用权重向量化二维数组列上的 bincount

问题描述

我一直在这里这里寻找解决方案,但没有看到如何将它应用到我的结构中。

我有 3 个数组:一个(M, N)零数组,一个(P,)索引(一些重复)和一个(P, N)值。

我可以用一个循环来完成它:

# a: (M, N)
# b: (P, N)
# ix: (M,)
for i in range(N):
    a[:, i] += np.bincount(ix, weights=b[:, i], minlength=M)

我还没有看到任何以这种方式或weights关键字使用索引的示例。我知道我需要将所有内容都放入一维数组中以对其进行矢量化,但是我正在努力弄清楚如何做到这一点。

标签: pythonnumpyvectorization

解决方案


基本思想与那些链接帖子中的一些详细讨论保持相同,即创建一个2Dbin 数组,每个要处理的“1D 数据”具有偏移量(在这种情况下是每个 col)。所以,考虑到这些,我们最终会得到这样的结果——

# Extent of bins per col
n = ix.max()+1

# 2D bins for per col processing
ix2D = ix[:,None] + n*np.arange(b.shape[1])

# Finally use bincount with those 2D bins as flattened and with
# flattened b as weights. Reshaping is needed to add back into "a".
a[:n] += np.bincount(ix2D.ravel(), weights=b.ravel(), minlength=n*N).reshape(N,-1).T

推荐阅读