如何在Pandas大数据集中高效筛选至少两列值相同的行对

针对大数据集（如500万行），原NumPy方案因生成n×n矩阵导致内存溢出，我们可以基于列组合分组的思路实现低内存消耗的解决方案，核心是避免全量行对比较，转而利用分组聚合找出共享列值组合的行。

一、筛选至少2列值相同的行对 #

思路：枚举所有列的两两组合，对每对列的值进行分组，每个分组内的行两两满足该两列值相同，收集所有这类行对并去重（仅保留i<j的有序对避免重复）。

import pandas as pd
import itertools

def find_pairs_with_at_least_n_common_cols(df, n=2):
    matching_pairs = set()
    # 生成所有长度为n的列组合
    col_combinations = itertools.combinations(df.columns, n)
    
    for cols in col_combinations:
        # 按当前列组合分组，获取每个组的行索引列表
        grouped_indices = df.groupby(list(cols)).groups
        for idx_list in grouped_indices.values():
            # 组内至少有2行才会产生有效行对
            if len(idx_list) >= 2:
                # 生成组内所有i<j的行对，加入集合去重
                for pair in itertools.combinations(idx_list, 2):
                    if pair[0] < pair[1]:
                        matching_pairs.add(pair)
                    else:
                        matching_pairs.add((pair[1], pair[0]))
    
    # 转换为列表或数组格式输出
    return list(matching_pairs)

# 测试示例DataFrame
df = pd.DataFrame({
    'A': [1, 2, 3, 4, 5],
    'B': [1, 1, 1, 1, 2],
    'C': [1, 1, 2, 3, 3],
    'D': [2, 7, 9, 8, 4]})

# 获取至少2列相同的行对
print(find_pairs_with_at_least_n_common_cols(df, n=2))
# 输出: [(0, 1)]

二、筛选至少3列值相同的行对 #

只需修改函数参数n=3即可，此时会枚举所有3列组合，找出在某3列值完全相同的行对：

# 获取至少3列相同的行对
print(find_pairs_with_at_least_n_common_cols(df, n=3))
# 输出: []（示例中无满足条件的行对）

方案优势 #

1. 内存友好：无需生成n×n的巨大矩阵，仅存储分组后的索引列表和最终行对，内存消耗与列数的组合数、符合条件的行对数量正相关，远低于全量比较方案。
2. 高效可扩展：时间复杂度为O(k^n * n + S)，其中k是列数，S是符合条件的行对总数。当列数不多时（如k≤10），即使处理500万行也能高效运行。
3. 灵活适配：通过修改n参数可直接切换筛选“至少2列”或“至少3列”的需求。

注意事项 #

• 如果数据集存在大量重复的列值组合，可考虑提前对DataFrame去重，减少分组处理的规模。
• 对于超大规模数据集（如1000万行以上），可结合并行计算库对分组过程进行加速。

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