计数索引以生成热图

时间:2019-03-25 16:02:56

标签: python numpy

我想累积指向m-by-n数组的索引,该索引指向具有该形状的另一个数组以生成热图。例如,这些索引:

[
    [0, 1, 2, 0, 1, 2]
    [0, 1, 0, 0, 0, 2]
]

将产生以下数组:

[
    [2, 0, 0]
    [1, 1, 0]
    [1, 0, 1]
]

我已经成功实现了一种算法,但是我开始怀疑,是否已经有针对此类问题的内置NumPy解决方案。

这是我的代码:

a = np.array([[0, 1, 2, 0, 1, 2], [0, 1, 0, 0, 0, 2]])

def _gather_indices(indices: np.ndarray, shape: tuple):
    heat = np.zeros(shape)
    for i in range(indices.shape[-1]):
        heat[tuple(indices[:, i])] += 1

1 个答案:

答案 0 :(得分:1)

可以建议两种方法。

使用np.add.at-

heat = np.zeros(shape,dtype=int)
np.add.at(heat,(a[0],a[1]),1)

或者使用基于tuple()的版本来获得更多美学之一-

np.add.at(heat,tuple(a),1)

使用bincount-

idx = np.ravel_multi_index(a,shape)
np.bincount(idx,minlength=np.prod(shape)).reshape(shape)

另外,我们可以使用shape中索引的最大限制来计算a-

shape = a.max(axis=1)+1

样品运行-

In [147]: a
Out[147]: 
array([[0, 1, 2, 0, 1, 2],
       [0, 1, 0, 0, 0, 2]])

In [148]: shape = (3,3)

In [149]: heat = np.zeros(shape,dtype=int)
     ...: np.add.at(heat,(a[0],a[1]),1)

In [151]: heat
Out[151]: 
array([[2, 0, 0],
       [1, 1, 0],
       [1, 0, 1]])

In [173]: idx = np.ravel_multi_index(a,shape)

In [174]: np.bincount(idx,minlength=np.prod(shape)).reshape(shape)
Out[174]: 
array([[2, 0, 0],
       [1, 1, 0],
       [1, 0, 1]])