在Pandas中高效计算分组变量间比率并保留原始数据

本文详细介绍了如何在pandas dataframe中，针对特定分组（如`g1`和`g2`），计算列`tpe`中`ts`与`td`类型对应的`qc`值之比，并将结果作为新行添加回原始dataframe。教程采用`set_index`、`unstack`和`div`的向量化方法，有效处理了缺失值，确保了计算效率和结果的准确性，避免了传统`groupby().apply()`方法的潜在性能瓶颈。

引言

在数据分析工作中，我们经常需要对DataFrame中的数据进行分组计算，特别是当数据以“堆叠”（stacked）形式存储时，例如将不同类型的数据作为同一列的不同行值。一个常见的需求是计算同一组内不同类型变量之间的比率，同时保留原始数据，并优雅地处理缺失值。本教程将展示一种高效且Pandas风格的解决方案，以解决此类问题。

问题描述与数据结构

假设我们有一个包含多列的DataFrame，其中G1和G2代表分组键，TPE列指示了数据的类型（例如'td'或'ts'），QC列则包含相应的值。我们的目标是：

1. 根据G1和G2进行分组。
2. 在每个分组内，计算TPE为'ts'的QC值与TPE为'td'的QC值之比（即ts/td）。
3. 将这个比率作为一个新的行添加到原始DataFrame中，新行的TPE列应标记为'ratio'。
4. 如果某个分组中缺少'td'或'ts'中的任何一个，则对应的比率应为空（NaN）。
5. 最终输出的DataFrame应包含所有原始行以及新计算出的比率行。

以下是示例输入DataFrame df_in：

import pandas as pd
import numpy as np

data = {
    'G1': ['A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'D'],
    'G2': ['S1', 'S1', 'S2', 'S2', 'S1', 'S1', 'S2', 'S2', 'S1', 'S2'],
    'TPE': ['td', 'ts', 'td', 'ts', 'td', 'ts', 'td', 'ts', 'td', 'ts'],
    'QC': [2, 4, 6, 3, 20, 40, 60, 30, 90, 7]
}
df_in = pd.DataFrame(data)

# 模拟缺失值情况
df_in.loc[df_in['G1'] == 'C', 'TPE'] = 'td' # 确保C只有td
df_in.loc[df_in['G1'] == 'D', 'TPE'] = 'ts' # 确保D只有ts
df_in.loc[df_in['G1'] == 'C', 'QC'] = 90
df_in.loc[df_in['G1'] == 'D', 'QC'] = 7

print("原始DataFrame df_in:")
print(df_in)

输出 df_in:

原始DataFrame df_in:
  G1  G2 TPE  QC
0  A  S1  td   2
1  A  S1  ts   4
2  A  S2  td   6
3  A  S2  ts   3
4  B  S1  td  20
5  B  S1  ts  40
6  B  S2  td  60
7  B  S2  ts  30
8  C  S1  td  90
9  D  S2  ts   7

解决方案：向量化方法

传统的groupby().apply()方法虽然灵活，但在处理大量数据时可能效率低下，尤其是在需要将结果重新组合回原始DataFrame时。更优的方案是利用Pandas的set_index、unstack和向量化操作。

步骤一：数据重塑 (set_index 和 unstack)

首先，我们将G1、G2和TPE设置为索引，然后对TPE进行unstack操作。这将把'td'和'ts'从行值转换为列名，并将QC列的值填充到相应的位置。

# 将G1, G2, TPE设置为索引，并对TPE进行unstack操作，获取QC列的值
tmp = df_in.set_index(['G1', 'G2', 'TPE']).unstack()['QC']
print("\n中间DataFrame tmp (unstack后):")
print(tmp)

输出 tmp:

中间DataFrame tmp (unstack后):
TPE   td    ts
G1 G2         
A  S1  2.0   4.0
   S2  6.0   3.0
B  S1 20.0  40.0
   S2 60.0  30.0
C  S1 90.0   NaN
D  S2  NaN   7.0

可以看到，unstack操作自动将缺失的'td'或'ts'值填充为NaN，这正是我们处理缺失比率所需要的。

步骤二：计算比率

在重塑后的tmp DataFrame中，'td'和'ts'已经成为了独立的列。现在，我们可以直接对这两列进行向量化除法操作来计算比率。

零一万物开放平台

零一万物大模型开放平台

下载

# 计算 ts/td 的比率
ratio_series = tmp['ts'].div(tmp['td'])
print("\n计算出的比率 Series:")
print(ratio_series)

输出 ratio_series:

计算出的比率 Series:
G1  G2
A   S1    2.0
    S2    0.5
B   S1    2.0
    S2    0.5
C   S1    NaN
D   S2    NaN
dtype: float64

这里，NaN值的传播是自动且正确的：如果td或ts中的任何一个为NaN，或者td为零，结果都将是NaN。

步骤三：重塑比率结果并添加TPE列

为了将比率结果与原始DataFrame合并，我们需要将其转换回与原始数据相似的格式。这包括将多级索引G1, G2重置为列，并将比率值命名为QC，同时添加一个TPE列并赋值为'ratio'。

# 将比率 Series 转换为 DataFrame，重置索引，并将比率列命名为'QC'
ratio_df = ratio_series.reset_index(name='QC')
# 添加TPE列，标记为'ratio'
ratio_df = ratio_df.assign(TPE='ratio')
print("\n比率DataFrame (待合并):")
print(ratio_df)

输出 ratio_df:

比率DataFrame (待合并):
  G1  G2   QC    TPE
0  A  S1  2.0  ratio
1  A  S2  0.5  ratio
2  B  S1  2.0  ratio
3  B  S2  0.5  ratio
4  C  S1  NaN  ratio
5  D  S2  NaN  ratio

步骤四：合并原始数据与比率数据

最后一步是使用pd.concat将原始DataFrame df_in与新生成的ratio_df合并。

# 合并原始DataFrame和比率DataFrame
df_out = pd.concat([df_in, ratio_df], ignore_index=True)
print("\n最终输出DataFrame df_out:")
print(df_out)

输出 df_out:

最终输出DataFrame df_out:
   G1  G2    TPE   QC
0   A  S1     td  2.0
1   A  S1     ts  4.0
2   A  S2     td  6.0
3   A  S2     ts  3.0
4   B  S1     td 20.0
5   B  S1     ts 40.0
6   B  S2     td 60.0
7   B  S2     ts 30.0
8   C  S1     td 90.0
9   D  S2     ts  7.0
10  A  S1  ratio  2.0
11  A  S2  ratio  0.5
12  B  S1  ratio  2.0
13  B  S2  ratio  0.5
14  C  S1  ratio  NaN
15  D  S2  ratio  NaN

完整代码示例

将上述步骤整合到一起，得到简洁高效的解决方案：

import pandas as pd
import numpy as np

# 原始DataFrame设置
data = {
    'G1': ['A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'D'],
    'G2': ['S1', 'S1', 'S2', 'S2', 'S1', 'S1', 'S2', 'S2', 'S1', 'S2'],
    'TPE': ['td', 'ts', 'td', 'ts', 'td', 'ts', 'td', 'ts', 'td', 'ts'],
    'QC': [2, 4, 6, 3, 20, 40, 60, 30, 90, 7]
}
df_in = pd.DataFrame(data)

# 模拟缺失值情况 (确保C只有td，D只有ts)
df_in.loc[df_in['G1'] == 'C', 'TPE'] = 'td'
df_in.loc[df_in['G1'] == 'D', 'TPE'] = 'ts'
df_in.loc[df_in['G1'] == 'C', 'QC'] = 90
df_in.loc[df_in['G1'] == 'D', 'QC'] = 7

# 解决方案核心代码
tmp = df_in.set_index(['G1', 'G2', 'TPE']).unstack()['QC']
ratio_df = tmp['ts'].div(tmp['td']).reset_index(name='QC').assign(TPE='ratio')
df_out = pd.concat([df_in, ratio_df], ignore_index=True)

print("最终结果 df_out:")
print(df_out)

注意事项与总结

1. 效率优势: 这种基于set_index().unstack().div()的向量化方法通常比groupby().apply()更高效，尤其是在处理大型数据集时，因为它避免了Python级别的循环。
2. 缺失值处理: unstack操作会自动为缺失的组合生成NaN，而div操作会正确地将NaN传播到比率结果中，无需额外的手动检查。
3. 数据类型: QC列的数据类型在计算比率后可能会变为浮点数，因为除法运算通常会产生浮点数。pd.concat会根据需要自动调整列的数据类型。
4. 可读性: 尽管链式操作可能看起来复杂，但一旦理解了set_index和unstack的工作原理，这种方法在Pandas中是非常直观和惯用的。

通过采用这种方法，我们能够以一种既高效又健壮的方式，在Pandas DataFrame中计算分组变量间的比率，并灵活地将结果集成回原始数据结构。