优化Pandas代码：计算分组中in_control为Y的站点占比

示例数据构造 #

先构造符合需求的示例DataFrame：

import pandas as pd

data = {
    'chr': ['chr1', 'chr1', 'chr2', 'chr2', 'chr2', 'chr3'],
    'start': [100, 200, 300, 400, 500, 600],
    'end': [150, 250, 350, 450, 550, 650],
    'in_control': ['Y', 'N', 'Y', 'Y', 'N', 'N'],
    'cluster': ['A', 'A', 'B', 'B', 'B', 'C']
}
df = pd.DataFrame(data)

现有繁琐代码示例（供对比） #

常规的分步处理方式通常是先分组统计再合并：

# 分组统计每个cluster的总站点数和控制站点数
cluster_stats = df.groupby('cluster').agg(
    total_sites=('cluster', 'size'),
    control_sites=('in_control', lambda x: (x == 'Y').sum())
)
# 计算比例并合并回原DataFrame
cluster_stats['cluster_sites_in_control'] = cluster_stats['control_sites'] / cluster_stats['total_sites']
df = df.merge(cluster_stats[['cluster_sites_in_control']], on='cluster', how='left')

简洁高效的解决方案 #

利用Pandas的groupby.transform方法，一行代码即可完成需求：

df['cluster_sites_in_control'] = df.groupby('cluster')['in_control'].transform(
    lambda x: (x == 'Y').mean()
)

原理说明 #

• groupby('cluster')：按cluster字段对数据分组
• transform：将分组计算的结果自动广播回原DataFrame对应组的每一行，保持原数据行数不变
• (x == 'Y').mean()：直接计算每组中in_control为'Y'的比例（均值等价于符合条件的数量除以组内总数量）

执行结果 #

运行简洁代码后，DataFrame新增列的结果如下：

优势对比 #

• 代码更简洁：仅一行核心代码完成需求，可读性更强
• 效率更高：避免了中间统计DataFrame的创建和合并操作，减少内存占用与计算步骤

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