Подсчитать индексы для массива, чтобы получить тепловую карту

Question

Я бы хотел накапливать индексы, которые указывают на массив 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

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