掩码基于索引的numpy数组

Question

如何根据实际索引值屏蔽数组？

也就是说,如果我有一个10 x 10 x 30的矩阵,我想在第一个和第二个索引相等时屏蔽数组.

例如,[1, 1 , :] 应该屏蔽,因为1和1相等,但[1, 2, :]不应该,因为它们不相同.

我只是问第三个维度,因为它类似于我当前的问题,可能会使事情复杂化.但我的主要问题是,如何根据索引的值掩盖数组？

Answer 1

通常,要访问索引的值,您可以使用np.meshgrid:

i, j, k = np.meshgrid(*map(np.arange, m.shape), indexing='ij')
m.mask = (i == j)

这种方法的优点是,它适用于任意布尔函数上i,j和k.它比使用identity特殊情况慢一点.

In [56]: %%timeit
   ....: i, j, k = np.meshgrid(*map(np.arange, m.shape), indexing='ij')
   ....: i == j
10000 loops, best of 3: 96.8 µs per loop

正如@Jaime指出的那样,meshgrid支持一个sparse选项,它没有那么多重复,但在某些情况下需要更多关注,因为它们不进行广播.它可以节省内存并加快速度.例如,

In [77]: x = np.arange(5)

In [78]: np.meshgrid(x, x)
Out[78]: 
[array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]]),
 array([[0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2],
       [3, 3, 3, 3, 3],
       [4, 4, 4, 4, 4]])]

In [79]: np.meshgrid(x, x, sparse=True)
Out[79]: 
[array([[0, 1, 2, 3, 4]]),
 array([[0],
       [1],
       [2],
       [3],
       [4]])]

所以,您可以sparse按照他的说法使用该版本,但您必须强制广播:

i, j, k = np.meshgrid(*map(np.arange, m.shape), indexing='ij', sparse=True)
m.mask = np.repeat(i==j, k.size, axis=2)

加速:

In [84]: %%timeit
   ....: i, j, k = np.meshgrid(*map(np.arange, m.shape), indexing='ij', sparse=True)
   ....: np.repeat(i==j, k.size, axis=2)
10000 loops, best of 3: 73.9 µs per loop