首页 > 解决方案 > 计算数组的索引以生成热图

问题描述

我想将指向一个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

标签: pythonnumpy

解决方案


可以提出两种方法。

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]])

推荐阅读