python - 从一维 NumPy 数组创建二维掩码

Question

我想从 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]]

我正在寻找不需要迭代的解决方案，因为结果数组将有数百万行。

Answer 1

在利用 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] = ....