根据1D计数器阵列填充2D阵列列

时间:2018-05-26 22:28:51

标签: python arrays numpy autofill

我正在寻找一个numpy解决方案来填充2D数组中的每一列("下面的例子中的#34;)和#34; 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.]]

使用numpy例程是否有更简单(更快)的方法来执行此操作而不需要每列循环?

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

像上面这样的东西会很好,但是不起作用。

感谢您的时间!

1 个答案:

答案 0 :(得分:2)

你可以使用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.]])