查找两个一维数组（有重复项）之间所有匹配值索引的最快方法

Question

问题描述

假设我们有两个简单的数组：

query = np.array([100, 4000, 500, 700, 400, 100])
match = np.array([6, 100, 4000, 100, 10, 8, 10])

我想在查询和匹配之间找到所有匹配值的索引。所以在这种情况下，结果将是：

value   query   match
100        0    1
100        0    3
100        5    1
100        5    3
4000       1    2

实际上这些数组将包含数百万个项目

---

“愚蠢的”循环解决方案

qs = []
query_locs = []
match_locs = []

for i in np.arange(query.size):
    q = query[i]
    # Get matching indexes in "match"
    match_loc = np.where(match == q)[0]
    n = match_loc.size
    # Update location arrays
    match_locs.extend(match_loc)
    query_locs.extend(np.repeat(i,n))
    # Store the matching value
    qs.extend(np.repeat(q,n))

result = np.vstack((qs, query_locs, match_locs)).T
print(result)
 [[ 100    0    1]
 [ 100    0    3]
 [4000    1    2]
 [ 100    5    1]
 [ 100    5    3]]

（也许numba 可以使这个循环非常快，但是当我尝试这个时，我得到了一些关于签名的错误，所以不确定）

---

Numpy 构建

有相当多的内置 numpy 函数可以解决此问题的唯一值，例如使用 searchsorted、intersect1d，但是，正如文档中所述，它们“返回已排序的唯一值”，因此不考虑到重复。 StackOverflow 上针对此问题的一些具有唯一值的示例：

• NumPy: Comparing Elements in Two Arrays
• Efficient way to compute intersecting values between two numpy arrays

我可以想象会有一种 更快 的方法来使用 numpy 而不是循环来做到这一点，所以很想看到答案！

Answer 1

您可以将一维数组转换为数据帧并进行连接，如下所示：

query = np.array([100, 4000, 500, 700, 400, 100])
match = np.array([6, 100, 4000, 100, 10, 8, 10])
dfquery = pd.DataFrame(range(len(query)), index=query, columns=['query'])
dfmatch = pd.DataFrame(range(len(match)), index=match, columns=['match'])
dfquery.join(dfmatch, how='inner')

结果：

    query   match
100     0       1
100     0       3
100     5       1
100     5       3
4000    1       2

Answer 2

你可以用 newaxis 破解它：

>>> comparison = np.equal(query[:, np.newaxis], match[np.newaxis, :])
array([[False,  True, False,  True, False, False, False],
       [False, False,  True, False, False, False, False],
       [False, False, False, False, False, False, False],
       [False, False, False, False, False, False, False],
       [False, False, False, False, False, False, False],
       [False,  True, False,  True, False, False, False]])

它本质上创建了笛卡尔积 (query x matches)（注意内存成本），然后应用二元函数 np.equal 将积空间中的每个元素有效地转换为 bool。 输出可以通过逐行读取来解释为：只要comparison[i, j] 是True，查询元素 i 就等于匹配元素 j。 您可以使用以下命令收集所有 True 对的索引：

list(zip(*comparison.nonzero()))
[(0, 1), (0, 3), (1, 2), (5, 1), (5, 3)]

ps：如果数组太长而无法构建产品，则按元素迭代它们是您唯一的选择。