使用Pandas groupby和transform实现复杂条件列填充

场景描述与需求分析

在数据分析工作中，我们经常需要根据复杂的业务逻辑为dataframe添加新列。考虑以下场景：我们有一个包含col1、col2和col3的dataframe。我们的目标是创建一个名为new_col的新列，其填充逻辑如下：

1. 按Col1进行分组。
2. 在每个分组内部，检查Col2列是否包含值'Y'。
   • 如果该分组内的Col2包含'Y'：那么该分组所有行的New_Col都应填充为对应Col2为'Y'那一行的Col3值。
   • 如果该分组内的Col2不包含'Y'：那么该分组所有行的New_Col都应直接复制其各自行的Col3值。

以下是原始数据示例：

我们期望得到的New_Col结果如下：

可以看到，对于Col1为1和3的分组，因为其中有Col2为'Y'的行，所以New_Col被填充为该行对应的Col3值'XX'。而对于Col1为2的分组，由于没有Col2为'Y'的行，New_Col直接复制了Col3的值。

Pandas核心操作解析

为了高效地实现上述需求，我们将组合使用Pandas的几个关键函数：

mask()函数：条件性替换值

DataFrame.mask(cond, other=nan, inplace=False, axis=None, level=None)mask()函数根据指定条件cond来替换DataFrame或Series中的值。如果cond为True，则替换为other（默认为NaN）；如果cond为False，则保留原值。这与where()函数的作用相反。在本例中，我们将用它来“隐藏”那些不符合我们首要条件的Col3值。

groupby()与transform()：分组操作与广播结果

DataFrame.groupby(by=None, axis=0, level=None, as_index=True, sort=True, group_keys=True, squeeze=NoDefault.no_default, observed=False, dropna=True)DataFrameGroupBy.transform(func, *args, **kwargs)groupby()用于将DataFrame按一个或多个列进行分组。transform()是GroupBy对象的一个方法，它对每个分组应用一个函数func，并返回一个与原始DataFrame或Series具有相同索引的Series，其结果会“广播”到整个分组。例如，transform('first')会返回每个分组的第一个非空值，并将其应用到该分组的所有行。

fillna()：处理缺失值

DataFrame.fillna(value=None, method=None, axis=None, inplace=False, limit=None, downcast=None)fillna()函数用于填充DataFrame或Series中的缺失值（NaN或None）。我们可以指定一个固定值、一个字典、一个Series，或者使用前向/后向填充方法。在本例中，我们将利用它来回填那些没有满足首要条件的行的Col3值。

解决方案步骤详解

我们将通过链式操作，分三步实现上述需求。

首先，初始化我们的DataFrame：

import pandas as pd
import numpy as np

data = {
    'Col1': [1, 1, 1, 2, 2, 3, 3, 3],
    'Col2': ['X', 'Y', 'X', 'X', 'X', 'X', 'X', 'Y'],
    'Col3': ['ABC', 'XX', 'QW', 'VB', 'AY', 'MM', 'YY', 'XX']
}
df = pd.DataFrame(data)
print("原始DataFrame:")
print(df)

输出:

原始DataFrame:
   Col1 Col2 Col3
0     1    X  ABC
1     1    Y   XX
2     1    X   QW
3     2    X   VB
4     2    X   AY
5     3    X   MM
6     3    X   YY
7     3    Y   XX

步骤1：隔离关键信息

我们首先需要识别出那些Col2为'Y'的行所对应的Col3值，并“隐藏”其他行的Col3值。这可以通过mask()函数实现。我们希望在Col2不等于'Y'时，将Col3的值替换为NaN。

Live PPT

一款AI智能化生成演示内容的在线工具。只需输入一句话、粘贴一段内容、或者导入文件，AI生成高质量PPT。

299

查看详情

# 隔离关键信息：当Col2不为'Y'时，将Col3替换为NaN
masked_col3 = df['Col3'].mask(df['Col2'] != 'Y')
print("\n步骤1：隔离关键信息 (masked_col3):")
print(masked_col3)

输出:

步骤1：隔离关键信息 (masked_col3):
0    NaN
1     XX
2    NaN
3    NaN
4    NaN
5    NaN
6    NaN
7     XX
Name: Col3, dtype: object

此时，我们得到了一个Series，其中只有Col2为'Y'的行的Col3值被保留，其余都被替换为NaN。

步骤2：分组提取并广播

接下来，我们需要根据Col1进行分组，并在每个分组内，找到步骤1中得到的Series的第一个非NaN值。如果分组内有'Y'，那么transform('first')会返回该'Y'对应的Col3值；如果分组内没有'Y'（即所有值都是NaN），transform('first')将返回None或NaN。这个结果会广播到整个分组。

# 分组提取并广播：获取每个分组的第一个非NaN值
grouped_transformed = masked_col3.groupby(df['Col1']).transform('first')
print("\n步骤2：分组提取并广播 (grouped_transformed):")
print(grouped_transformed)

输出:

步骤2：分组提取并广播 (grouped_transformed):
0      XX
1      XX
2      XX
3    None
4    None
5      XX
6      XX
7      XX
Name: Col3, dtype: object

可以看到，对于Col1为1和3的分组，所有行都得到了'XX'。而对于Col1为2的分组，因为原始masked_col3中对应这些行的值都是NaN，所以transform('first')也返回了None（或NaN）。

步骤3：回填默认值

最后一步，我们需要处理那些在步骤2中得到None/NaN的行。根据需求，这些行应该填充为它们各自原始的Col3值。这可以通过fillna()函数实现。

# 回填默认值：将None/NaN值替换为原始的Col3值
df['New_Col'] = grouped_transformed.fillna(df['Col3'])
print("\n步骤3：回填默认值并生成最终DataFrame:")
print(df)

输出:

步骤3：回填默认值并生成最终DataFrame:
   Col1 Col2 Col3 New_Col
0     1    X  ABC      XX
1     1    Y   XX      XX
2     1    X   QW      XX
3     2    X   VB      VB
4     2    X   AY      AY
5     3    X   MM      XX
6     3    X   YY      XX
7     3    Y   XX      XX

至此，我们成功地创建了New_Col，并按照预期的逻辑进行了填充。

完整代码实现

将上述所有步骤组合成一个简洁的链式操作，是Pandas数据处理的常见模式：

import pandas as pd
import numpy as np

# 原始数据
data = {
    'Col1': [1, 1, 1, 2, 2, 3, 3, 3],
    'Col2': ['X', 'Y', 'X', 'X', 'X', 'X', 'X', 'Y'],
    'Col3': ['ABC', 'XX', 'QW', 'VB', 'AY', 'MM', 'YY', 'XX']
}
df = pd.DataFrame(data)

# 使用链式操作创建新列
df['New_Col'] = (df['Col3']
                 .mask(df['Col2'] != 'Y')  # 步骤1: 隔离关键信息
                 .groupby(df['Col1'])      # 步骤2: 分组
                 .transform('first')       # 步骤2: 提取并广播
                 .fillna(df['Col3']))      # 步骤3: 回填默认值

print("最终DataFrame:")
print(df)

注意事项

1. 数据类型： transform('first')在处理不同数据类型时，对于全NaN的分组，返回的结果可能是None（对于object类型）或NaN（对于数值类型）。fillna()函数能够很好地处理这两种情况。
2. 性能： 这种链式操作通常比使用循环或apply()函数更高效，尤其是在处理大型数据集时，因为它充分利用了Pandas底层的优化。
3. 通用性： 这种模式非常灵活，可以根据不同的条件和聚合函数（如'max', 'min', lambda函数等）来修改mask()和transform()的参数，以适应各种复杂的条件填充需求。例如，如果需求是取分组内Col2为'Y'的最后一个Col3值，可以将transform('first')改为transform('last')。

总结

本文介绍了一种高效且优雅的Pandas数据处理技巧，通过巧妙地组合mask()、groupby()、transform()和fillna()函数，实现了根据分组内条件动态填充新列的需求。这种方法不仅代码简洁，而且在处理大规模数据时表现出优秀的性能。掌握这种链式操作模式，将有助于您更灵活、更高效地进行数据清洗和特征工程。

以上就是使用Pandas groupby和transform实现复杂条件列填充的详细内容，更多请关注php中文网其它相关文章！