基于分组和条件添加新列的 Pandas 教程

本文介绍了如何使用 Pandas 在 DataFrame 中基于分组和条件计算并添加新列。我们将通过一个实际案例，演示如何根据 ‘text’ 列进行分组，并根据 ‘number’ 列的值动态计算 ‘test’ 列的值，其中’number’列的值会影响计算的步长。

问题描述

假设我们有一个 DataFrame，包含 ‘id’, ‘date’, ‘date_difference’, ‘number’ 和 ‘text’ 等列。我们的目标是创建一个名为 ‘test’ 的新列，其值取决于 ‘text’ 列的分组以及 ‘number’ 列的值。具体规则如下：

1. 根据 ‘text’ 列进行分组。
2. 在每个分组内，’date’ 列按降序排列。
3. 当 ‘number’ 列的值为 0 时，步长从 1 开始。
4. 如果在分组内找到 ‘number’ 列的值为 1，则步长增加 1。
5. 如果分组内没有 ‘number’ 列的值为 1，则整个分组的步长保持为 1。

解决方案

我们可以使用 Pandas 的 groupby()、apply()、sort_values()、shift() 和 cumsum() 等方法来实现这个目标。

代码示例

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)  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”): 这会将 DataFrame 按照 ‘text’ 列的值进行分组。
3. .apply(lambda g: …): apply 方法将一个函数应用到每个分组（这里用 g 表示每个分组后的 DataFrame）。
4. g.sort_values(by=”date”, ascending=False): 在每个分组内，按照 ‘date’ 列降序排序。
5. .number.shift(periods=1, fill_value=1): 将 ‘number’ 列的值向上移动一位（即向前 shift），并将第一个缺失值（由于 shift 造成的）填充为 1。 shift 操作的目的是将当前行的 number 值与上一行的 number 值关联起来，从而确定步长。fill_value=1 确保了第一个值的步长至少为 1。
6. .cumsum(): 计算移动后的 ‘number’ 列的累积和。这个累积和就是我们想要的 ‘test’ 列的值。
7. .droplevel(“text”): 移除由 groupby 引入的索引层级 “text”，使结果与原始 DataFrame 的索引对齐。

输出结果

运行上述代码，将得到以下 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

可以看到，’test’ 列已经按照我们的规则正确计算出来了。

注意事项

• 确保 ‘date’ 列的数据类型是 datetime 类型。如果不是，可以使用 pd.to_datetime() 函数进行转换。
• fill_value 的值需要根据实际情况进行调整。在本例中，我们将其设置为 1，以确保步长至少为 1。
• 理解 shift 函数的作用至关重要。它将当前行的 ‘number’ 值与上一行的 ‘number’ 值关联起来，从而实现了动态步长的计算。

总结

本文演示了如何使用 Pandas 在 DataFrame 中基于分组和条件计算并添加新列。通过 groupby()、apply()、sort_values()、shift() 和 cumsum() 等方法的组合，我们可以灵活地处理各种复杂的数据计算任务。 掌握这些技巧可以帮助你更高效地进行数据分析和处理。