Python 对比两列数据提取公共元素及对应行号的实现方法

Python没有直接实现该需求的内置函数，但可以通过标准库语法快速实现，也可以借助numpy库的现成方法完成，以下是两种可直接运行的实现：

方案1：纯Python实现（无需第三方库） #

适合数据量较小的场景，可灵活控制结果排序规则：

A = [11, 24, 7, 7, 0, 4, 9, 20, 3, 5]
B = [5, 8, 9, 1, 11]

# 构造A的元素到索引的映射
a_index_map = {val: idx for idx, val in enumerate(A)}

common_vals = []
row_a = []
row_b = []

# 按B的遍历顺序收集结果，匹配你给出的示例输出顺序
for idx_b, val in enumerate(B):
    if val in a_index_map:
        common_vals.append(val)
        row_a.append(a_index_map[val])
        row_b.append(idx_b)

print(common_vals) # 输出：[5, 9, 11]
print(row_a) # 输出：[9, 6, 0]
print(row_b) # 输出：[0, 2, 4]

如果需要结果顺序和公共元素在A中出现的顺序一致，只需调换遍历逻辑，先构造B的元素索引映射，再遍历A即可。

方案2：numpy实现（适合大数据量场景） #

可以直接调用numpy.intersect1d方法，开启return_indices参数即可同时返回公共元素在两个数组中的索引：

import numpy as np

A = np.array([11, 24, 7, 7, 0, 4, 9, 20, 3, 5])
B = np.array([5, 8, 9, 1, 11])

# 直接获取公共元素、对应A的索引、对应B的索引
common_vals, row_a, row_b = np.intersect1d(A, B, return_indices=True)

print(common_vals) # 输出：[ 5  9 11]
print(row_a) # 输出：[9 6 0]
print(row_b) # 输出：[0 2 4]

注意：numpy.intersect1d返回的公共元素默认按升序排序，如果需要自定义顺序可以自行调整结果排序规则。如果列表存在重复公共元素，该方法默认只返回每个元素第一次出现的索引。

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