python - 用于屏蔽每列单个切片的矢量化方法

Question

我有一个 numpy 数组:

>>> a = np.arange(20).reshape(5, -1)
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15],
       [16, 17, 18, 19]])

我有一个按列顺序排列的区域数组，我想为它们创建一个 bool 掩码:

idx = np.array([[0,2], [1,3], [2,4], [1,4]])

我想要的这组索引掩码是:

array([[ True, False, False, False],
       [ True,  True, False,  True],
       [False,  True,  True,  True],
       [False, False,  True,  True],
       [False, False, False, False]])

因此 0 列有 0:2 屏蔽，1 列有 1:3 屏蔽，等等。我目前的方法有效，但我正在寻找矢量化的东西:

def foo(a, idx):
    out = np.zeros(a, dtype=np.bool8)
    for (i, j), k in zip(idx, np.arange(a[1])):
        out[i:j, k] = True
    return out

在行动中:

foo(a.shape, idx)

array([[ True, False, False, False],
       [ True,  True, False,  True],
       [False,  True,  True,  True],
       [False, False,  True,  True],
       [False, False, False, False]])

Answer 1

使用 broadcasting -

In [434]: r = np.arange(a.shape[0])[:,None]

In [435]: (idx[:,0] <= r) & (idx[:,1] > r)
Out[435]: 
array([[ True, False, False, False],
       [ True,  True, False,  True],
       [False,  True,  True,  True],
       [False, False,  True,  True],
       [False, False, False, False]])