從2D NumPy數組中查找1D數組？

Question

假設我有

[[array([x1, y1]), z1]
 [array([x2, y1]), z2]
 ......
 [array([xn, yn]), zn]
]

我想找到array([x5, y5])的索引。 如何使用NumPy有效地查找？

Answer 1

首先，由於混合的數據格式，我認為您不能以矢量化方式提取數組。 因此，您可以使用loop comprehension從每個列表元素中提取與數組對應的第一個元素作為2D數組。 因此，假設A是輸入列表，我們將-

arr = np.vstack([a[0] for a in A])

然后，只需使用NumPy's broadcasting feature以矢量化方式進行比較，因為它將沿所有行廣播該比較，並使用np.all(axis=1)查找所有匹配的行。 最后，使用np.flatnonzero獲取最終索引。 因此，這個難題的最終和平將是-

idx = np.flatnonzero((arr == search1D).all(1))

您可以閱讀this post的答案，以查看其他替代方法，以獲取在2D數組問題中進行此類1D數組搜索的索引。

樣品運行-

In [140]: A
Out[140]: 
[[array([3, 4]), 11],
 [array([2, 1]), 12],
 [array([4, 2]), 16],
 [array([2, 1]), 21]]

In [141]: search1D = [2,1]

In [142]: arr = np.vstack([a[0] for a in A]) # Extract 2D array

In [143]: arr
Out[143]: 
array([[3, 4],
       [2, 1],
       [4, 2],
       [2, 1]])

In [144]: np.flatnonzero((arr == search1D).all(1)) # Finally get indices
Out[144]: array([1, 3])