统计DataFrame组合出现次数的循环逻辑错误，如何修正？

问题描述 #

我有两个包含组合数据的DataFrame（均含20000+行），需要遍历df1并统计其中的组合在df2中的出现次数。例如，df1中的组合"B,D"应在df2的"A,B,D"、"B,C,D"、"A,B,C,D"这3个组合中被识别（因为这些组合包含"B,D"的所有元素）。

现有代码 #

import pandas as pd

d1 = {'Combinations' : ['A,B', 'A,C', 'A,D', 'B,C', 'B,D', 'C,D', 'A,B,C', 'A,C,D', 'B,C,D', 'A,B,D', 'A,B,C,D',]}
df1 = pd.DataFrame(data=d1)

d2 = {'Combinations' : ['A,B,D', 'A,C,D', 'B,C,D', 'A,B,C,D', 'A,D']}
df2 = pd.DataFrame(data=d2)

counts = []

for i, row in df1.iterrows():
    values = row['Combinations'].split(',')

    count = df2.Combinations.str.contains('|'.join(values)).sum()
    counts.append(count)
    
df1['counts'] = counts

当前错误输出 #

Combinations  counts
0          A,B       5
1          A,C       5
2          A,D       5
3          B,C       4
4          B,D       5
5          C,D       5
6        A,B,C       5
7        A,C,D       5
8        B,C,D       5
9        A,B,D       5
10     A,B,C,D       5

期望正确输出 #

Combinations  counts
0          A,B       2
1          A,C       2
2          A,D       4
3          B,C       2
4          B,D       3
5          C,D       3
6        A,B,C       1
7        A,C,D       2
8        B,C,D       2
9        A,B,D       2
10     A,B,C,D       1

解决方案 #

原代码错误在于使用str.contains('|'.join(values))，这里的|是逻辑或关系，只要组合包含任意一个元素就会被计数，而我们需要的是同时包含所有元素的逻辑与。

修改后的代码 #

import pandas as pd

d1 = {'Combinations' : ['A,B', 'A,C', 'A,D', 'B,C', 'B,D', 'C,D', 'A,B,C', 'A,C,D', 'B,C,D', 'A,B,D', 'A,B,C,D',]}
df1 = pd.DataFrame(data=d1)

d2 = {'Combinations' : ['A,B,D', 'A,C,D', 'B,C,D', 'A,B,C,D', 'A,D']}
df2 = pd.DataFrame(data=d2)

counts = []

for _, row in df1.iterrows():
    # 将当前组合拆分为元素集合
    target_set = set(row['Combinations'].split(','))
    # 统计df2中包含所有目标元素的组合数量
    count = sum(
        target_set.issubset(set(comb.split(','))) 
        for comb in df2['Combinations']
    )
    counts.append(count)
    
df1['counts'] = counts
print(df1)

代码说明 #

1. 将df1的每个组合转换为集合，利用集合的issubset方法判断该集合是否是df2中某个组合集合的子集（即包含所有元素）。
2. 遍历df2的所有组合，逐个判断并累加符合条件的数量。

针对大数据集的优化版本

提前将df2的所有组合转换为集合列表，避免循环内重复拆分操作，提升性能：

import pandas as pd

d1 = {'Combinations' : ['A,B', 'A,C', 'A,D', 'B,C', 'B,D', 'C,D', 'A,B,C', 'A,C,D', 'B,C,D', 'A,B,D', 'A,B,C,D',]}
df1 = pd.DataFrame(data=d1)

d2 = {'Combinations' : ['A,B,D', 'A,C,D', 'B,C,D', 'A,B,C,D', 'A,D']}
df2 = pd.DataFrame(data=d2)

# 提前处理df2的组合为集合，减少重复计算
df2_sets = [set(comb.split(',')) for comb in df2['Combinations']]

counts = []
for _, row in df1.iterrows():
    target_set = set(row['Combinations'].split(','))
    count = sum(target_set.issubset(s) for s in df2_sets)
    counts.append(count)
    
df1['counts'] = counts
print(df1)

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