优化基于元组对的DataFrame概率计算逻辑（替代低效apply传参）

原始数据与需求 #

现有包含字符串配对的DataFrame：

import pandas as pd

d = {'value': [['Red', 'Blue'],['Blue', 'Yellow'],['Blue', 'Yellow'],['Yellow', 'Orange'],['Green', 'Purple'],['Purple', 'Yellow'],['Yellow', 'Red'],['Violet', 'Blue'],['Blue', 'Green'],['Green', 'Red'],['Red', 'Brown'],['Blue', 'Green']]}
df = pd.DataFrame(data=d)

需求：为每行计算概率，规则为当前行配对的出现行数 ÷ 以该配对首元素开头的总行数。

原始低效实现 #

原方案通过apply逐行调用函数，每次都全表扫描统计次数，效率极低（时间复杂度O(n²)）：

def find_prob(df, tup):
    d = df[df.value.apply(lambda x: x[0] == tup[0] and x[1] == tup[1])].shape[0]
    p = df[df.value.apply(lambda x: x[0] == tup[0])].shape[0]
    return d / p

df['probs'] = df.value.apply(lambda x: find_prob(df, x))

高效优化方案 #

利用pandas的分组统计能力，仅需几次批量操作即可完成，时间复杂度大幅降低：

步骤说明 #

1. 将value列拆分为start和end两列，简化分组操作；
2. 批量统计每个(start, end)配对的出现次数；
3. 批量统计每个start的总出现次数；
4. 计算每个配对的概率并映射回原DataFrame。

完整代码 #

import pandas as pd

# 1. 拆分配对为两列
df[['start', 'end']] = pd.DataFrame(df['value'].tolist(), index=df.index)

# 2. 统计每个(start, end)的出现次数
pair_counts = df.groupby(['start', 'end']).size().rename('pair_count')

# 3. 统计每个start的总出现次数
start_counts = df.groupby('start').size().rename('start_count')

# 4. 计算每个配对的概率
prob_map = (pair_counts / start_counts).reset_index(name='probs')

# 5. 将概率映射回原DataFrame
df = df.merge(prob_map, on=['start', 'end'], how='left')

# 可选：如果不需要start/end列，可以删除
# df = df.drop(['start', 'end'], axis=1)

效果验证 #

执行后df['probs']列即为所需概率，例如：

• 配对['Blue', 'Yellow']出现2次，以Blue开头的总行数是4，概率为2/4=0.5；
• 配对['Blue', 'Green']出现2次，概率同样为2/4=0.5。

为什么更高效 #

• 原方案每行都要遍历全表两次，数据量越大越慢；
• 优化方案仅通过两次分组统计完成批量计算，后续通过合并/映射完成赋值，避免了重复扫描，数据量越大，效率提升越明显。

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