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])]

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

Numpy find identical elements from list of arrays and another array

python

arrays

identity

numpy