使用numpy.where（）或类似方法从矩阵中的一行获取特定值

Question

As an MWE, I have a 2d numpy array:

import numpy as np
n1 = np.array([[6,7,8,1], [5,2,4,8], [3,4,2,1], [8,7,2,10]])
n1
array([[ 6,  7,  8,  1],
       [ 5,  2,  4,  8],
       [ 3,  4,  2,  1],
       [ 8,  7,  2, 10]])

I want to get the indices where '6' and '8' occur in the first row; where '2' and '8' appear in the second row; where '3' and '4' appear in the third row, and where '7' and '2' appear in the last row. That is, I have a list of numpy arrays:

list1 = [np.array([6, 8]), np.array([2, 8]), np.array([3, 4]), np.array([7, 2])]

I want [n1[i, np.where(ar1)] for i in range(len(list1))] or something like that to return a new list with the column where the value from list1 occurs:

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

Clearly I've tried [n1[i, np.where(ar1)] for i in range(len(list1))] . Any ideas?

Answer 1

Here's one using NumPy broadcasting for array output -

In [117]: arr_list1 = np.array(list1)

In [118]: mask = (n1[:,:,None] == arr_list1[:,None,:]).any(2)

In [119]: np.where(mask)[1].reshape(-1,2)
Out[119]: 
array([[0, 2],
       [1, 3],
       [0, 1],
       [1, 2]])

Explanation

• Basically we are extending both n1 and arr_list1 to 3D arrays by introducing singleton dims/dims with lengths 1 with None/np.newaxis , such that when compared against each other would result in elementwise comparison against the last axes from those two arrays as a full 3D array.

• We then look for ANY match along the last axis which has a length of 2 corresponding to the two elems in arr_list1 for each row. That gives us a 2D array. Finally, we need the matching row indices, so np.where()[1] .

Step-by-step run for a closer look -

1) Inputs :

In [124]: n1
Out[124]: 
array([[ 6,  7,  8,  1],
       [ 5,  2,  4,  8],
       [ 3,  4,  2,  1],
       [ 8,  7,  2, 10]])

In [125]: arr_list1
Out[125]: 
array([[6, 8],
       [2, 8],
       [3, 4],
       [7, 2]])

2) Comparison :

In [126]: (n1[:,:,None] == arr_list1[:,None,:])
Out[126]: 
array([[[ True, False],
        [False, False],
        [False,  True],
        [False, False]],

       [[False, False],
        [ True, False],
        [False, False],
        [False,  True]],

       [[ True, False],
        [False,  True],
        [False, False],
        [False, False]],

       [[False, False],
        [ True, False],
        [False,  True],
        [False, False]]], dtype=bool)

3) ANY reduction :

In [127]: (n1[:,:,None] == arr_list1[:,None,:]).any(2)
Out[127]: 
array([[ True, False,  True, False],
       [False,  True, False,  True],
       [ True,  True, False, False],
       [False,  True,  True, False]], dtype=bool)

In [128]: mask = _

4) Finally get matching row indices :

In [130]: np.where(mask)[1].reshape(-1,2)
Out[130]: 
array([[0, 2],
       [1, 3],
       [0, 1],
       [1, 2]])