如何匹配两个numpy数组中包含的值对

Question

我有两组坐标，想找出coo 集合中的哪些坐标与targets 集合中的任何坐标相同。我想知道coo 集合中的索引，这意味着我想获取索引或布尔值列表。

import numpy as np

coo = np.array([[1,2],[1,6],[5,3],[3,6]]) # coordinates
targets = np.array([[5,3],[1,6]]) # coordinates of targets

print(np.isin(coo,targets))

[[ True False]
 [ True  True]
 [ True  True]
 [ True  True]]

期望的结果是以下两种之一：

[False True True False] # bool list
[1,2] # list of concerning indices

我的问题是……

• np.isin 没有 axis-属性，所以我可以使用 axis=1。
• 即使对输出的每一行应用 logical and 也会为最后一个元素返回 True，这是错误的。

我知道循环和条件，但我确信 Python 提供了更优雅的解决方案。

Answer 1

对于大型数组，此解决方案的扩展性会更差，在这种情况下，其他建议的答案会表现得更好。

---

这是利用broadcasting的一种方法：

(coo[:,None] == targets).all(2).any(1)
# array([False,  True,  True, False])

---

详情

检查coo 中的每一行是否与target 中的另一行匹配，方法是直接比较已将第一个轴添加到coo，因此它可以针对targets 进行广播：

(coo[:,None] == targets)

array([[[False, False],
        [ True, False]],

       [[False, False],
        [ True,  True]],

       [[ True,  True],
        [False, False]],

       [[False, False],
        [False,  True]]])

然后检查第二个轴上哪个ndarrays 具有all 值到True：

(coo[:,None] == targets).all(2)

array([[False, False],
       [False,  True],
       [ True, False],
       [False, False]])

最后使用any 检查哪些行至少有一个True。

Answer 2

这是一个简单直观的解决方案，它实际上使用numpy.isin()，来匹配元组，而不是匹配单个数字：

# View as a 1d array of tuples
coo_view     = coo.view(dtype='i,i').reshape((-1,))
targets_view = targets.view(dtype='i,i').reshape((-1,))

result = np.isin(coo_view, targets_view)
print (result)
print(result.nonzero()[0])

输出：

[False  True  True False]
[1 2]

注意事项：

1. 这些views 的创建不涉及任何数据复制。
2. dtype='i,i' 指定我们希望视图的每个元素都是两个整数的元组

Answer 3

numpy_indexed 包以矢量化方式实现了这种类型的功能（免责声明：我是它的作者）。遗憾的是，numpy 缺少很多开箱即用的功能。我开始 numpy_indexed 的目的是将它合并到 numpy 中，但是存在一些向后兼容性问题，并且像这样的大包往往会移动缓慢。所以这在过去 3 年没有发生过；但是现在 python 打包生态系统运行得非常好，只需在您的环境中再添加一个包就很简单了。

import numpy_indexed as npi
bools = npi.in_(targets, coo)

这将具有类似于@fountainhead 发布的解决方案的时间复杂度（根据当前接受的答案，是对数而不是线性），但 npi 库也会为您提供自动化测试的安全性，以及如果您决定从稍微不同的角度解决问题，还有很多其他方便的选择。