Numpy find identical elements from list of arrays and another array-CodePudding

Suppose I have two lists a and b, a is 1D array with sub-arrays, b is 2D array without sub-arrays. How can I find the elements in a identical with b with keeping the structure of sub arrays in a?

   a = [array([[1,1,1],
               [2,2,2],
               [3,4,5]]),
        array([[0,0,0],
               [3,3,3],
               [5,5,5],
               [9,3,3],
               [8,2,2]]),
        ]

   b = array([[0,1,2],
              [3,4,5],
              [5,5,5],
              [9,3,3],
              [9,9,9]])

In this case, the identical elements are [3,4,5],[5,5,5],[9,3,3] I want the result to be like: (as well as extracting the index in a)

 >>> [array([[3,4,5]]),
      array([[5,5,5],
             [9,3,3]])
      ]

Is there any easy way to do it in numpy？

CodePudding user response：

Assuming a is a list of arrays, you can use broadcasting to perform the comparisons of all elements:

out = [x[(x == b[:,None]).all(2).any(0)] for x in a]

Output:

[array([[3, 4, 5]]),
 array([[5, 5, 5],
        [9, 3, 3]])]

Indices:

[np.where((x == b[:,None]).all(2).any(0))[0]  for x in a]

Output:

[array([2]), array([2, 3])]