如何用非重复代码聚合Pandas DataFrame？替代多筛选与合并操作

我现在需要汇总一个DataFrame，现有代码能得到正确结果，但重复度极高，想避免这种情况。我觉得可以用groupby、agg、apply这类方法，但还没找到可行的实现方式。目标是生成最终的df_summ，现在我用了太多筛选行的中间DataFrame，还有多次merge来整合结果，肯定有更简洁的方法，但一直没理清。

实际输入的df_stats有数百万行，输出的df_summ有几十列，以下是最小可复现代码：

import io
import pandas as pd

TESTDATA="""
enzyme  regions   N   length
AaaI    all       10  238045
AaaI    all       20  170393
AaaI    all       30  131782
AaaI    all       40  103790
AaaI    all       50  81246
AaaI    all       60  62469
AaaI    all       70  46080
AaaI    all       80  31340
AaaI    all       90  17188
AaaI    captured  10  292735
AaaI    captured  20  229824
AaaI    captured  30  193605
AaaI    captured  40  163710
AaaI    captured  50  138271
AaaI    captured  60  116122
AaaI    captured  70  95615
AaaI    captured  80  73317
AaaI    captured  90  50316
AagI    all       10  88337
AagI    all       20  19144
AagI    all       30  11030
AagI    all       40  8093
AagI    all       50  6394
AagI    all       60  4991
AagI    all       70  3813
AagI    all       80  2759
AagI    all       90  1666
AagI    captured  10  34463
AagI    captured  20  19220
AagI    captured  30  15389
AagI    captured  40  12818
AagI    captured  50  10923
AagI    captured  60  9261
AagI    captured  70  7753
AagI    captured  80  6201
AagI    captured  90  4495
"""

df_stats = pd.read_csv(io.StringIO(TESTDATA), sep='\\s+')

df_cap_N90 = df_stats[(df_stats['N'] == 90) & (df_stats['regions'] == 'captured')].drop(columns=['regions', 'N'])
df_cap_N50 = df_stats[(df_stats['N'] == 50) & (df_stats['regions'] == 'captured')].drop(columns=['regions', 'N'])

df_all_N50 = df_stats[(df_stats['N'] == 50) & (df_stats['regions'] == 'all')     ].drop(columns=['regions', 'N'])

df_summ_cap_N50_all_N50 = pd.merge(df_cap_N50, df_all_N50, on='enzyme', how='inner', suffixes=('_cap_N50', '_all_N50'))
df_summ_cap_N50_all_N50['cap_N50_all_N50'] = (df_summ_cap_N50_all_N50['length_cap_N50'] -
                                              df_summ_cap_N50_all_N50['length_all_N50'])
print(df_summ_cap_N50_all_N50)

df_summ_cap_N90_all_N50 = pd.merge(df_cap_N90, df_all_N50, on='enzyme', how='inner', suffixes=('_cap_N90', '_all_N50'))
df_summ_cap_N90_all_N50['cap_N90_all_N50'] = df_summ_cap_N90_all_N50['length_cap_N90'] - df_summ_cap_N90_all_N50['length_all_N50']
print(df_summ_cap_N90_all_N50)

df_summ = pd.merge(df_summ_cap_N50_all_N50.drop(columns=['length_cap_N50', 'length_all_N50']),
                   df_summ_cap_N90_all_N50.drop(columns=['length_cap_N90', 'length_all_N50']),
                   on='enzyme', how='inner')
print(df_summ)

原代码输出：

enzyme  length_cap_N50  length_all_N50  cap_N50_all_N50
0   AaaI          138271           81246            57025
1   AagI           10923            6394             4529
  enzyme  length_cap_N90  length_all_N50  cap_N90_all_N50
0   AaaI           50316           81246           -30930
1   AagI            4495            6394            -1899
  enzyme  cap_N50_all_N50  cap_N90_all_N50
0   AaaI            57025           -30930
1   AagI             4529            -1899

生物信息学背景说明
（可跳过）
这段代码是生物信息学项目中的一步，目的是根据限制性内切酶切割DNA的特性筛选最优酶。输入表格里，enzyme列是酶名称，regions列是两种DNA区域类型，N列是DNA切割精细程度的统计量（类似分位数N10到N90），length是对应统计量的核苷酸长度值。需要计算诸如cap_N90_all_N50 = captured区域N90长度 - all区域N50长度这类指标，再用这些指标给酶排序。

解决方案：用Pivot重塑数据，减少重复操作 #

核心思路是先把长格式的df_stats转成宽格式，让每个酶对应的所有统计值都在同一行，这样后续计算指标会非常方便，不需要多次筛选和merge。

优化后的代码 #

import io
import pandas as pd

TESTDATA="""
enzyme  regions   N   length
AaaI    all       10  238045
AaaI    all       20  170393
AaaI    all       30  131782
AaaI    all       40  103790
AaaI    all       50  81246
AaaI    all       60  62469
AaaI    all       70  46080
AaaI    all       80  31340
AaaI    all       90  17188
AaaI    captured  10  292735
AaaI    captured  20  229824
AaaI    captured  30  193605
AaaI    captured  40  163710
AaaI    captured  50  138271
AaaI    captured  60  116122
AaaI    captured  70  95615
AaaI    captured  80  73317
AaaI    captured  90  50316
AagI    all       10  88337
AagI    all       20  19144
AagI    all       30  11030
AagI    all       40  8093
AagI    all       50  6394
AagI    all       60  4991
AagI    all       70  3813
AagI    all       80  2759
AagI    all       90  1666
AagI    captured  10  34463
AagI    captured  20  19220
AagI    captured  30  15389
AagI    captured  40  12818
AagI    captured  50  10923
AagI    captured  60  9261
AagI    captured  70  7753
AagI    captured  80  6201
AagI    captured  90  4495
"""

df_stats = pd.read_csv(io.StringIO(TESTDATA), sep='\\s+')

# 1. 重塑数据：将regions和N合并为列名，每个酶一行
df_pivot = df_stats.pivot(
    index='enzyme',
    columns=['regions', 'N'],
    values='length'
).reset_index()

# 2. 给列名重命名，改成类似length_cap_N50的格式
df_pivot.columns = [
    'enzyme' if col[0] == 'enzyme' 
    else f'length_{col[0]}_N{col[1]}' 
    for col in df_pivot.columns
]

# 3. 计算需要的指标
df_pivot['cap_N50_all_N50'] = df_pivot['length_captured_N50'] - df_pivot['length_all_N50']
df_pivot['cap_N90_all_N50'] = df_pivot['length_captured_N90'] - df_pivot['length_all_N50']

# 4. 提取最终需要的列，得到df_summ
df_summ = df_pivot[['enzyme', 'cap_N50_all_N50', 'cap_N90_all_N50']]

print(df_summ)

输出结果 #

enzyme  cap_N50_all_N50  cap_N90_all_N50
0   AaaI            57025           -30930
1   AagI             4529            -1899

优势说明 #

• 减少重复代码：不管要计算多少个指标，只需要一次pivot操作，后续直接列运算即可，不用反复筛选和merge
• 性能更优：对于百万行的大数据，pivot的效率远高于多次筛选+merge，因为只需要一次数据重组
• 扩展性强：如果后续需要增加新的指标（比如cap_N10_all_N30），只需要加一行列运算代码即可，不需要新增筛选和merge步骤

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