[英]Filling of 2D array columns according to 1D counter array
我正在寻找一种numpy解决方案,以在二维数组(在下面的示例中为“ a”)中的每一列填充不同的一维计数器数组(在以下示例中为“ cnt”)中定义的多个“ 1”值。
我尝试了以下方法:
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)
像上面的东西会很好,但是不起作用。
感谢您的时间!
您可以使用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.]])
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.