Pandas数据透视与向量化操作：高效聚合复杂数据集

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'])

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_N90['length_cap_N90'] - df_summ_cap_N90_all_N90['length_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("原始方法计算结果：")
print(df_summ)

# 1. 筛选相关数据
# 仅保留N为50或90的行，因为只有这些N值参与最终计算
filtered_df = df_stats.loc[df_stats["N"].isin([50, 90])]

# 2. 使用pivot进行数据透视
# index='enzyme'：以enzyme作为新的行索引
# columns=['regions', 'N']：将regions和N的组合作为新的列索引（多级列索引）
# values='length'：透视后单元格的值取自length列
pivoted_df = filtered_df.pivot(index="enzyme", columns=["regions", "N"], values="length")
print("\n透视后的DataFrame (pivoted_df):")
print(pivoted_df)

# 3. 执行向量化计算
# 提取'captured'区域的N50和N90长度
captured_lengths = pivoted_df["captured"]
# 提取'all'区域的N50长度
all_N50_length = pivoted_df[("all", 50)]

# 计算 (captured N50 - all N50) 和 (captured N90 - all N50)
# captured_lengths.sub(all_N50_length, axis=0)
# axis=0 表示按行进行广播，即captured_lengths的每一行都减去all_N50_length的对应行值
result_df = captured_lengths.sub(all_N50_length, axis=0)

# 4. 调整列名并重置索引
# 为结果列添加前缀和后缀，使其符合目标输出格式
# add_prefix("cap_N")：为'captured'下的N值（50, 90）添加前缀'cap_N'
# add_suffix("_all_N50")：为所有结果列添加后缀'_all_N50'
final_output = result_df.add_prefix("cap_N").add_suffix("_all_N50").reset_index()

print("\n最终优化后的计算结果 (final_output):")
print(final_output)