Заполнение столбцов двумерного массива в соответствии с массивом одномерного счетчика

Question

Я ищу решение для кучи для заполнения каждого столбца в двумерном массиве («a» в приведенном ниже примере) количеством значений «1», как определено в другом массиве 1D-счетчиков («cnt» впример ниже).

Я пробовал следующее:

import numpy as np

cnt = np.array([1, 3, 2, 4])   # cnt array: how much elements per column are 1
a = np.zeros((5, 4))           # array that must be filled with 1s per column
for i in range(4):             # for each column
    a[:cnt[i], i] = 1          # all elements from top to cnt value are filled

print(a)

и дает желаемый результат:

[[1. 1. 1. 1.] 
 [0. 1. 1. 1.]
 [0. 1. 0. 1.]
 [0. 0. 0. 1.]
 [0. 0. 0. 0.]]

Есть ли более простой (и более быстрый) способ срутинная подпрограмма, чтобы сделать это без цикла для каждого столбца?

a = np.full((5, 4), 1, cnt)

Что-то вроде вышеизложенного было бы неплохо, но не работает.

Спасибо за ваше время!

Answer 1

Вы можете использовать np.where и вещать так:

>>> import numpy as np
>>> 
>>> cnt = np.array([1, 3, 2, 4])   # cnt array: how much elements per column are 1
>>> a = np.zeros((5, 4))           # array that must be filled with 1s per column
>>> 
>>> res = np.where(np.arange(a.shape[0])[:, None] < cnt, 1, a)
>>> res
array([[1., 1., 1., 1.],
       [0., 1., 1., 1.],
       [0., 1., 0., 1.],
       [0., 0., 0., 1.],
       [0., 0., 0., 0.]])

Или на месте:

>>> a[np.arange(a.shape[0])[:, None] < cnt] = 1
>>> a
array([[1., 1., 1., 1.],
       [0., 1., 1., 1.],
       [0., 1., 0., 1.],
       [0., 0., 0., 1.],
       [0., 0., 0., 0.]])