如何用Pandas双重循环统计子Reddit间的共享评论者

问题背景 #

你有一个包含subreddit列和对应评论作者列表列的Pandas DataFrame，想要遍历每一个subreddit，再和它之后的所有subreddit配对，统计两者的共享评论者数量，最终生成包含subreddit_1、subreddit_2和shared_users列的结果DataFrame。你尝试用双重循环但遇到了困难，不清楚怎么用apply实现。

原始数据示例 #

你的数据结构类似如下形式：

subreddit user
0xProject [7878ayush, Mr_Yukon_C, NomChompsky92, PM_ME_Y...
100sexiest [T10rock]
100yearsago [PM_ME_MII, Quisnam]
...
zyzz [Xayv]

输入输出示例 #

输入测试DataFrame：

import pandas as pd
df = pd.DataFrame({
    'subreddit': ['sub1', 'sub2', 'sub3', 'sub4'], 
    'user': [['A', 'B', 'C'], ['A', 'F', 'C'], ['F', 'E', 'D'], ['X', 'Y', 'Z']]
})

预期针对sub1的输出片段：

subreddit_1 subreddit_2 shared_users
sub1        sub2        2
sub1        sub3        0
sub1        sub4        0

解决方案 #

咱们不用低效的双重循环，而是用更高效的集合操作和组合生成来实现，步骤如下：

1. 将用户列表转换为集合 #

集合的交集操作比列表高效得多，先把每个subreddit的用户列转成集合，方便后续快速计算共享用户：

df['user_set'] = df['user'].apply(set)

2. 生成所有需要配对的subreddit组合 #

使用itertools.combinations来生成所有i<j的subreddit索引对，这样就能保证每个配对只计算一次（比如sub1和sub2，不会重复计算sub2和sub1），避免冗余工作：

from itertools import combinations

# 生成所有不重复的索引配对
index_pairs = combinations(df.index, 2)

3. 计算每对的共享用户数并构建结果 #

遍历这些索引对，取出对应的subreddit名称和用户集合，计算两个集合交集的长度（也就是共享用户的数量），然后收集所有结果：

results = []
for idx1, idx2 in index_pairs:
    sub1 = df.loc[idx1, 'subreddit']
    sub2 = df.loc[idx2, 'subreddit']
    shared_count = len(df.loc[idx1, 'user_set'] & df.loc[idx2, 'user_set'])
    results.append({
        'subreddit_1': sub1,
        'subreddit_2': sub2,
        'shared_users': shared_count
    })

# 将结果转换为目标DataFrame
result_df = pd.DataFrame(results)

完整代码示例 #

把上面的步骤整合起来，针对测试输入的完整可运行代码：

import pandas as pd
from itertools import combinations

# 测试输入DataFrame
df = pd.DataFrame({
    'subreddit': ['sub1', 'sub2', 'sub3', 'sub4'], 
    'user': [['A', 'B', 'C'], ['A', 'F', 'C'], ['F', 'E', 'D'], ['X', 'Y', 'Z']]
})

# 转换用户列表为集合
df['user_set'] = df['user'].apply(set)

# 生成索引组合并计算共享用户数
index_pairs = combinations(df.index, 2)
results = []
for idx1, idx2 in index_pairs:
    results.append({
        'subreddit_1': df.loc[idx1, 'subreddit'],
        'subreddit_2': df.loc[idx2, 'subreddit'],
        'shared_users': len(df.loc[idx1, 'user_set'] & df.loc[idx2, 'user_set'])
    })

result_df = pd.DataFrame(results)
print(result_df)

运行后输出的完整结果：

subreddit_1 subreddit_2  shared_users
0        sub1        sub2             2
1        sub1        sub3             0
2        sub1        sub4             0
3        sub2        sub3             1
4        sub2        sub4             0
5        sub3        sub4             0

为什么不用双重循环？ #

你最初的双重循环思路虽然可行，但效率很低，尤其是当DataFrame行数很多的时候。用itertools.combinations可以避免重复计算，同时集合的交集操作比列表遍历匹配要快得多，能大幅提升处理速度。

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