如何用Pandas/Numpy高效提取共享两值的行对应的无序值对？

问题描述 #

给定每行值不重复的DataFrame，需找出任意两行中存在至少两个共同值的情况（值所在列顺序无关），收集这些共同值对。每组值对最多对应一对行。

示例输入 #

d = {'col1': [1, 2, 5, 9], 'col2': [2, 7, 1, 2], 'col3': [3, 3, 2, 7]}
df = pd.DataFrame(data=d)

对应的DataFrame：

预期输出 #

1  2
2  3
2  7

现有实现（不够简洁） #

用户提供的Numpy循环实现：

def func(x):
    rows, columns = x.shape[0], x.shape[1]
    res = []
    for i in range(0,rows):
        for j in range(i+1, rows):
            aux = np.intersect1d(x[i,:], x[j,:])
            if aux.size>1:
                res.append(aux)
    return res

调用func(df.values)输出：

[array([2, 3]), array([1, 2]), array([2, 7])]

优化方案 #

方案1：Pandas + itertools 简洁实现 #

利用集合快速求交集，代码更易读简洁：

import pandas as pd
import itertools

def find_common_pairs(df):
    # 将每行转换为集合，便于快速计算交集
    row_sets = df.apply(set, axis=1).tolist()
    result = []
    # 遍历所有不重复的行对组合
    for i, j in itertools.combinations(range(len(row_sets)), 2):
        common_values = row_sets[i] & row_sets[j]
        if len(common_values) >= 2:
            # 排序后加入结果，保证输出格式一致
            result.append(sorted(common_values))
    # 转换为DataFrame返回，匹配预期输出格式
    return pd.DataFrame(result)

调用示例：

print(find_common_pairs(df))

输出：

0  1
0  2  3
1  1  2
2  2  7

方案2：Numpy向量化优化（适合大规模数据） #

通过向量化操作减少显式循环，提升处理效率：

import numpy as np
import pandas as pd

def find_common_pairs_np(df):
    arr = df.values
    # 生成所有上三角的行对索引（避免重复行对）
    i, j = np.triu_indices(arr.shape[0], k=1)
    
    # 将所有值映射为唯一索引，构建one-hot矩阵
    unique_vals = np.unique(arr)
    val_idx_map = {v: idx for idx, v in enumerate(unique_vals)}
    one_hot = np.zeros((arr.shape[0], len(unique_vals)), dtype=bool)
    for row_idx, row in enumerate(arr):
        one_hot[row_idx, [val_idx_map[v] for v in row]] = True
    
    # 计算每对行的交集数量，筛选出交集≥2的行对
    intersection_counts = (one_hot[i] & one_hot[j]).sum(axis=1)
    valid_pair_mask = intersection_counts >= 2
    
    # 提取有效行对的交集并排序
    result = []
    for row_i, row_j in zip(arr[i[valid_pair_mask]], arr[j[valid_pair_mask]]):
        common = np.intersect1d(row_i, row_j)
        result.append(np.sort(common))
    
    return pd.DataFrame(result)

调用后同样得到符合预期的输出。

方案说明 #

• 方案1代码简洁直观，适合中小规模数据集，可读性强。
• 方案2采用向量化处理，减少Python层循环，在数据量较大时性能更优。

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