本文详细介绍了如何使用 Pandas 在 DataFrame 中基于分组和条件添加新列。通过 groupby()、apply()、sort_values()、shift() 和 cumsum() 等函数的组合使用，可以实现复杂的数据转换和列生成。本文提供清晰的代码示例和详细的步骤解释，帮助读者理解并掌握该技巧，从而解决实际数据处理中的类似问题。
在数据分析和处理中，经常需要在 Pandas DataFrame 中基于现有列进行分组，并根据分组内的条件计算新列的值。本教程将详细介绍如何使用 Pandas 实现这一目标，并通过一个具体的例子进行说明。
示例数据
首先，我们创建一个示例 DataFrame，用于演示如何添加新列：
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 列进行分组，并根据 number 列的值生成一个新的列 test。对于每个分组，test 列的值的计算规则如下：
- 首先按照 date 列降序排列组内数据。
- 如果 number 列的值为 0，则 test 列的值从 1 开始递增。
- 如果在 number 列中出现 1，则后续的 test 列的值在之前的最大值基础上保持不变。
- 如果分组内没有 1，则 test 列的值始终为 1。
解决方案
以下代码展示了如何使用 Pandas 实现上述逻辑：
import pandas as pd
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)
out = df.assign(
# We assign the following values to the series name "test"
test=df
# Group on "text" -- if we grouped on ["text", "number"] we wouldn't see different numbers within the groups.
.groupby("text")
# Apply a chain of methods to the group (a pd.DataFrame).
.apply(
lambda g: (
# We sort "date" in descending order as you mention this partially controls the step size.
g.sort_values(by="date", ascending=False)
# We shift "number" forward one period with a fill_value of 1 for any newly introduced nulls.
.number.shift(periods=1, fill_value=1)
# Cumulatively sum the shifted "number" values
.cumsum()
)
# This will result in the new series, albeit sorted by descending "date".
)
# Drop the "text" level of the new multi-index.
.droplevel("text")
# The assign method acts as join, rearranging the newly created series to match the index of `df`.
)
print(out)代码解释：
- df.groupby("text"): 首先，我们使用 groupby() 函数按照 text 列对 DataFrame 进行分组。
- .apply(lambda g: ...): 然后，我们使用 apply() 函数对每个分组应用一个自定义的 lambda 函数。
- g.sort_values(by="date", ascending=False): 在 lambda 函数内部，我们首先按照 date 列降序排列每个分组。
- .number.shift(periods=1, fill_value=1): 使用 shift() 函数将 number 列的值向前移动一位。fill_value=1 用于填充移动后产生的缺失值，保证第一行数据的值为 1。
- .cumsum(): 使用 cumsum() 函数计算移动后的 number 列的累积和，从而得到 test 列的值。
- .droplevel("text"): 删除分组后产生的MultiIndex，确保结果与原DataFrame的索引对齐。
- df.assign(test=...): 最后，使用 assign() 函数将新生成的 test 列添加到原始 DataFrame 中。
总结
本教程介绍了如何使用 Pandas 在 DataFrame 中基于分组和条件添加新列。通过 groupby()、apply()、sort_values()、shift() 和 cumsum() 等函数的组合使用，可以灵活地实现复杂的数据转换和列生成。理解并掌握这些技巧，可以帮助您更高效地进行数据分析和处理。
注意事项：
- 在实际应用中，需要根据具体的数据和需求调整代码。
- 理解每个函数的具体作用和参数，才能更好地应用它们。
- 可以使用其他 Pandas 函数和方法来进一步优化代码，例如 fillna()、where() 等。
