numpy获取行和列索引2D数组的所有组合

Question

我有一个2D numpy数组如下:

import numpy as np
foo = np.array([[(i+1)*(j+1) for i in range(10)] for j in range(5)])

    #array([[ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10],
    #       [ 2,  4,  6,  8, 10, 12, 14, 16, 18, 20],
    #       [ 3,  6,  9, 12, 15, 18, 21, 24, 27, 30],
    #       [ 4,  8, 12, 16, 20, 24, 28, 32, 36, 40],
    #       [ 5, 10, 15, 20, 25, 30, 35, 40, 45, 50]])

我使用np.nonzero创建了一些过滤条件:

csum = np.sum(foo,axis=0)
#array([ 15,  30,  45,  60,  75,  90, 105, 120, 135, 150])
rsum = np.sum(foo,axis=1)
#array([ 55, 110, 165, 220, 275])
cfilter = np.nonzero(csum > 80)
#(array([5, 6, 7, 8, 9]),)
rfilter = np.nonzero(rsum < 165)
#(array([0, 1]),)

现在是否有一些优雅的numpy切片方法来获得rfilter中的r和cfilter中的r的所有foo [r,c]组合？即我想获得以下输出:

array([[ 6,  7,  8,  9, 10],
       [12, 14, 16, 18, 20]])

注意:我知道很容易进行基本切片选择以从数组中获取块,但在更高级的用例中,cfilter和rfilter中的索引不一定紧挨着彼此.

非常感谢!

Answer 1

要索引交叉产品,请使用np.ix_:

foo[np.ix_(*(rfilter + cfilter))]

您可以直接使用布尔索引(即不使用np.nonzero):

foo[np.ix_(np.sum(foo, axis=1) < 165, np.sum(foo, axis=0) > 80)]

请注意,所有np.ix_操作都是适当地添加轴以提供可以一起广播的索引数组:

>>> np.ix_(*(rfilter + cfilter))
(array([[0],
       [1]]), array([[5, 6, 7, 8, 9]]))