我想从 numpy 数组的行中获取前 n 个值,其中 n 在单独的一维数组中指定:
import numpy as np
a = np.zeros((5, 5))
n = [1, 3, 2, 4, 1]
result = [[1, 0, 0, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 0, 0, 0],
[1, 1, 1, 1, 0],
[1, 0, 0, 0, 0]]
我正在寻找不需要迭代的解决方案,因为结果数组将有数百万行。
最佳答案
在利用 broadcasting
时,对范围数组使用 n
的外部比较创建掩码,从而创建最终数组 -
ncols = 5
mask_out = np.greater.outer(n,np.arange(ncols))
sample 运行-
In [19]: n = [1, 3, 2, 4, 1]
In [9]: ncols = 5
# Output as mask
In [10]: np.greater.outer(n,np.arange(ncols))
Out[10]:
array([[ True, False, False, False, False],
[ True, True, True, False, False],
[ True, True, False, False, False],
[ True, True, True, True, False],
[ True, False, False, False, False]])
# Output as array of 0s and 1s
In [11]: np.greater.outer(n,np.arange(ncols)).view('i1')
Out[11]:
array([[1, 0, 0, 0, 0],
[1, 1, 1, 0, 0],
[1, 1, 0, 0, 0],
[1, 1, 1, 1, 0],
[1, 0, 0, 0, 0]], dtype=int8)
如果你必须填充一个已经初始化的数组 result
,只需使用 mask_out
来屏蔽数组,即 result[mask_out] = ...
.
关于python - 从一维 NumPy 数组创建二维掩码,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/55190295/