Numpy 从数组列表和另一个数组中找到相同的元素

Question

假设我有两个列表a和b ， a是带有子数组的一维数组， b 是没有子数组的二维数组。 如何在a中找到与b相同的元素，同时将子数组的结构保持在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]])

在这种情况下，相同的元素是[3,4,5],[5,5,5],[9,3,3]我希望结果如下：（以及在a中提取索引）

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

在numpy中有什么简单的方法吗？

Answer 1

假设a是一个数组列表，您可以使用广播来执行所有元素的比较：

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

输出：

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

指数：

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

输出：

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