Выберите элементы (n, n, 2) numpy массива с (n, n) -образной маской в ​​numpy без использования петель

Question

У меня есть массив (n, n, 2) numpy, элементы которого я хочу выбрать на основе маски (n, n) без использования циклов. Есть ли способ векторизовать эту операцию в numpy? Скажем, у меня есть numpy массив

X = array([[[18,  8],
        [ 9,  2],
        [11,  4],
        [18, 14]],

       [[ 8, 10],
        [13,  5],
        [13,  6],
        [13, 18]],

       [[ 8,  4],
        [ 2, 13],
        [19, 11],
        [ 3, 15]],

       [[12,  6],
        [ 7,  3],
        [19, 17],
        [ 1, 12]]])

и маска

M = array([[1, 0, 0, 0],
           [1, 1, 0, 0],
           [0, 0, 1, 0],
           [0, 0, 0, 0]])

Обрабатывая каждую двумерную запись в X как один элемент, есть ли способ использовать маску M выбрать элементы X? То есть, выберите 2-D элемент в X, если его соответствующий элемент в маске M равен 1.

Так что приведенный выше пример вернет

[
 [[18,  8]],
 [[ 8, 10],
  [13,  5]],
 [[19, 11]],
 []
]

Answer 1

Ну, вы можете просто использовать: печать (X [M == 1])

Answer 2

In [393]: I,J = np.nonzero(M)                                                                  
In [394]: I,J                                                                                  
Out[394]: (array([0, 1, 1, 2]), array([0, 0, 1, 2]))
In [395]: X[I,J,:]                                                                             
Out[395]: 
array([[18,  8],
       [ 8, 10],
       [13,  5],
       [19, 11]])

Было отмечено, что вы хотите не просто эти элементы, вы хотите, чтобы они группировались по строкам. Чтобы сделать это, понимание списка работает просто:

In [407]: [x[m.astype(bool)] for x,m in zip(X,M)]                                              
Out[407]: 
[array([[18,  8]]), array([[ 8, 10],
        [13,  5]]), array([[19, 11]]), array([], shape=(0, 2), dtype=int64)]
In [408]: [x[m.astype(bool)].tolist() for x,m in zip(X,M)]                                     
Out[408]: [[[18, 8]], [[8, 10], [13, 5]], [[19, 11]], []]

Да, использует al oop, но это то, что вы должны ожидать при получении списка элементов переменного размера.

* Решения 1007 *, no-l oop работают только с результатами обычного массива.