这个numpy数组索引的缩短版本

Question

我在python(numpy数组或scipy.sparse.matrices)中有以下代码,它有效:

X[a,:][:,b]

但它看起来并不优雅.'a'和'b'是一维布尔掩码.

'a'与X.shape [0]具有相同的长度,'b'与X.shape [1]具有相同的长度

我试过X[a,b]但它不起作用.

我想要完成的是同时选择特定的行和列.例如,选择第0,7,8行,然后从该结果中选择第2,3,4列中的所有行

你会如何让它更短更优雅？

Answer 1

你可以使用np.ix_这样的broadcasted indexing,像这样 -

X[np.ix_(a,b)]

虽然这不会比原始代码短,但希望应该更快.这是因为我们正在避免中间输出,就像使用X[a,:]一个切片创建的原始代码一样,然后是另一个切片X[a,:][:,b]来为我们提供最终输出.

此外,此方法适用于a和b作为数组int和boolean数组.

样品运行

In [141]: X = np.random.randint(0,99,(6,5))

In [142]: m,n = X.shape

In [143]: a = np.in1d(np.arange(m),np.random.randint(0,m,(m)))

In [144]: b = np.in1d(np.arange(n),np.random.randint(0,n,(n)))

In [145]: X[a,:][:,b]
Out[145]: 
array([[17, 81, 64],
       [87, 16, 54],
       [98, 22, 11],
       [26, 54, 64]])

In [146]: X[np.ix_(a,b)]
Out[146]: 
array([[17, 81, 64],
       [87, 16, 54],
       [98, 22, 11],
       [26, 54, 64]])

运行时测试

In [147]: X = np.random.randint(0,99,(600,500))

In [148]: m,n = X.shape

In [149]: a = np.in1d(np.arange(m),np.random.randint(0,m,(m)))

In [150]: b = np.in1d(np.arange(n),np.random.randint(0,n,(n)))

In [151]: %timeit X[a,:][:,b]
1000 loops, best of 3: 1.74 ms per loop

In [152]: %timeit X[np.ix_(a,b)]
1000 loops, best of 3: 1.24 ms per loop