Блог

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

Можно было бы предложить два метода.

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