本文旨在讲解如何使用 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 列，其值取决于以下规则：
- 组内按 date 降序排列。
- 如果 number 列的值为 0，则步长从 1 开始。
- 如果在组内找到 number 列的值为 1，则步长增加 1。
- 如果组内没有 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)
代码解释
-
df.assign(test=…): 使用 assign 函数创建一个新的列 test，其值将由后续的计算得出。
-
df.groupby(“text”): 按照 text 列进行分组。这是实现按组计算的关键步骤。
-
.apply(lambda g: …): 对每个分组应用一个 lambda 函数。这个 lambda 函数接收一个 DataFrame g 作为参数，代表一个分组的数据。
-
g.sort_values(by=”date”, ascending=False): 在每个分组内，按照 date 列降序排列。
-
.number.shift(periods=1, fill_value=1): 将 number 列的值向上移动一位。fill_value=1 表示对于移动后产生的缺失值，用 1 填充。 这一步是为了实现步长的累加。
-
.cumsum(): 对移动后的 number 列进行累加求和。这会根据条件增加步长。
-
.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，代表一个分组的数据。
- 根据实际需求调整排序、移动和累加的参数。
- 在处理大规模数据时，注意性能优化，避免不必要的计算。
