高效查找两个含重复值的大型数组共享值的配对索引

以下述简单数组为例：

# 索引位置0,1,2,3,4,5
a = np.array([1,1,3,4,6])
b = np.array([6,6,1,3])

需求为从两个数组中获取所有匹配值对应的双方索引配对：例如值1对应配对0,2和1,2，完整输出如下：

0,2 # 匹配值1
1,2 # 匹配值1
2,3 # 匹配值3
4,0 # 匹配值6
4,1 # 匹配值6

注意：两个数组均未排序，且包含重复元素——这是其他同类解答常默认不满足的两个前提条件。上述示例规模很小，但实际需要将方案应用于单组长度约4万的数组场景。

1. Python双层循环方案 #

indx = []
for i, aval in enumerate(a):
    for j, bval in enumerate(b):
        if aval == bval:
            indx.append([i,j])
# 输出结果 [[0, 2], [1, 2], [2, 3], [4, 0], [4, 1]]

2. Python字典方案 #

from collections import defaultdict
adict = defaultdict(list)
bdict = defaultdict(list)
for i, aval in enumerate(a): adict[aval].append(i)
for j, bval in enumerate(b): bdict[bval].append(j)

for val, a_positions in adict.items():
    for b_position in bdict[val]:
        for a_position in a_positions:
            print(a_position, b_position)

3. Numpy where方案 #

print(np.where(a.reshape(-1,1) ==  b))

4. Polars DataFrame方案 #

将数组转换为DataFrame后使用Polars实现：

import polars as pl
a_df = pl.DataFrame( {'x': a, 'apos':list(range(len(a)))} )
b_df = pl.DataFrame( {'x': b, 'bpos':list(range(len(b)))} )
result = a_df.join(b_df, how='inner', on='x')

大数据量场景测试 #

4万长度数组的测试场景下，Polars方案目前速度最快，耗时约0.02秒。创建DataFrame再做连接的方式比自研方案速度表现更好，因此好奇是否存在性能更优的实现方式。

测试数据生成代码：

import numpy as np
a = np.random.randint(0,1000, 40000)
b = np.random.randint(0,1000, 40000)

以上述测试数据运行各方案的耗时统计如下：

1. Python双层循环：218s
2. Python字典方案：0.03s
3. numpy.where：4.5s
4. Polars连接方案：0.02s

• 多数同类问题仅返回单个数组内的匹配值索引，不同时返回双方索引
• 部分同类问题分别返回A、B数组的匹配索引，但不返回两两配对的索引组合（可参考前文示例输出格式）

目前DataFrame类库的实现速度最优，希望了解是否存在其他方案可以超越该速度，Cython、numba、pythran等任意技术栈的实现均可接受。

内容的提问来源于stack exchange，提问作者CodeNoob