Question

我试图找到最有效的方法来获取另一个数组中嵌套数组的索引。

import numpy as np
#                     0     1      2     3
haystack = np.array([[1,3],[3,4,],[5,6],[7,8]])
needles  = np.array([[3,4],[7,8]])

鉴于needles中包含的数组，我想在haystack中找到它们的索引。在这种情况下1,3。

我想出了这个解决方案：

 indexes = [idx for idx,elem in enumerate(haystack) if elem in needles ]

这是错误的，因为实际上elem中的一个元素位于needles中足以返回idx。

有没有更快的选择？

Answer 1

此响应提供了一个类似问题Get intersecting rows across two 2D numpy arrays的解决方案，您使用了np.in1d函数，该函数非常有效，但是您可以通过给出两个数组的视图来实现，从而可以处理它们作为一维数据数组。就您而言，您可以

A = np.array([[1,3],[3,4,],[5,6],[7,8]])
B = np.array([[3,4],[7,8]])
nrows, ncols = A.shape
dtype={'names':['f{}'.format(i) for i in range(ncols)],
       'formats':ncols * [A.dtype]}
indexes, = np.where(np.in1d(A.view(dtype), B.view(dtype)))

输出：

print(indexes)
> array([1, 3])

Answer 2

您可以尝试

indices = np.apply_along_axis(lambda x: np.where(np.isin(haystack, x).sum(axis=1)==2)[0], 1, needles).flatten()
indices
>>> array([1, 3])

Answer 3

仅在嵌套的最上层搜索时，您可以直接使用索引功能。

indexes = [[a,haystack.index(a)] for a in needles.tolist()]

编辑：仅留下此答案以提供替代方法，但使用核心numpy函数（例如在其他答案中找到的那些函数）可能是更好的方法

如何找到嵌套在另一个数组中的数组的索引？

3 个答案: