Pandas教程：分组与条件判断添加列方法

本文旨在讲解如何使用 Pandas 在数据框中基于分组和条件判断来创建新的列。通过 groupby()、apply()、sort_values()、shift() 和 cumsum() 等函数，可以实现复杂的数据转换和计算，从而生成符合特定业务逻辑的新列。文章提供详细的代码示例和步骤解释，帮助读者理解并掌握该技巧。

Pandas 提供了强大的数据处理功能，其中基于分组和条件判断创建新列是一项常见的任务。本教程将介绍如何利用 Pandas 的 groupby()、apply() 以及其他相关函数，根据特定条件和分组规则生成新的列。

示例数据

首先，我们创建一个示例 DataFrame，它包含了id、date、date_difference、number 和 text 等列。我们的目标是基于 text 列进行分组，并根据 number 列的值以及日期顺序，生成一个新的 test 列。

import pandas as pd
import numpy as np

data = {
    'id': [1, 2, 3, 4, 5, 6, 7],
    'date': ['2019-02-01', '2019-02-10', '2019-02-25', '2019-03-05', '2019-03-16', '2019-04-05', '2019-05-15'],
    'date_difference': [None, 9, 15, 11, 10, 19, 40],
    'number': [1, 0, 1, 0, 0, 0, 0],
    'text': ['A', 'A', 'A', 'A', 'A', 'B', 'B']
}

df = pd.DataFrame(data)
print(df)

目标

对于每个 text 组，我们希望创建一个 test 列，其值取决于以下规则：

1. 组内按 date 降序排列。
2. 如果 number 列的值为 0，则步长从 1 开始。
3. 如果在组内找到 number 列的值为 1，则步长增加 1。
4. 如果组内没有 number 列的值为 1，则步长保持为 1。

解决方案

以下代码展示了如何使用 Pandas 实现上述目标：

out = df.assign(
    test=df
    .groupby("text")
    .apply(
        lambda g: (
            g.sort_values(by="date", ascending=False)
            .number.shift(periods=1, fill_value=1)
            .cumsum()
        )
    )
    .droplevel("text")
)
print(out)

代码解释

1. df.assign(test=...): 使用 assign 函数创建一个新的列 test，其值将由后续的计算得出。

2. df.groupby("text"): 按照 text 列进行分组。这是实现按组计算的关键步骤。

3. .apply(lambda g: ...): 对每个分组应用一个 lambda 函数。这个 lambda 函数接收一个 DataFrame g 作为参数，代表一个分组的数据。

4. g.sort_values(by="date", ascending=False): 在每个分组内，按照 date 列降序排列。

5. .number.shift(periods=1, fill_value=1): 将 number 列的值向上移动一位。fill_value=1 表示对于移动后产生的缺失值，用 1 填充。 这一步是为了实现步长的累加。

6. .cumsum(): 对移动后的 number 列进行累加求和。这会根据条件增加步长。

7. .droplevel("text"): groupby 操作会引入一个额外的层级，使用 droplevel 函数移除该层级，使得结果的索引与原始 DataFrame 匹配。

输出结果

运行上述代码，将得到以下结果：

   id        date  date_difference  number text  test
0   1  2019-02-01              NaN       1    A     2
1   2  2019-02-10              9.0       0    A     2
2   3  2019-02-25             15.0       1    A     1
3   4  2019-03-05             11.0       0    A     1
4   5  2019-03-16             10.0       0    A     1
5   6  2019-04-05             19.0       0    B     1
6   7  2019-05-15             40.0       0    B     1

总结

本教程展示了如何使用 Pandas 的 groupby() 和 apply() 函数，结合 sort_values()、shift() 和 cumsum() 函数，基于分组和条件判断创建新的列。这种方法可以灵活地处理各种复杂的数据转换和计算任务。理解并掌握这些技巧，可以有效地提高数据处理的效率和准确性。

注意事项

• 确保理解分组的含义和目标，选择合适的分组列。
• 在 apply() 函数中，注意 lambda 函数的输入参数是 DataFrame，代表一个分组的数据。
• 根据实际需求调整排序、移动和累加的参数。
• 在处理大规模数据时，注意性能优化，避免不必要的计算。