pandas 按分组将两列上移动增加新列

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

pandas 中的 shift() 方法，可以对一个 Series（或者说 DataFrame 中的一列）上下做位移，在本例中，我们会操作若干个列同时位移，而且这个位移还是按分组进行的。

注：本例在 pandas 2.2 下通过，同时建议安装 pyarrow 库。

数据与需求

我们构造数据如下：

import pandas as pd
from io import StringIO

data = '''
a,b,c
1,4,2
1,3,4
2,6,6
2,4,0
2,7,3
'''

df = pd.read_csv(StringIO(data))
df
'''
   a  b  c
0  1  4  2
1  1  3  4
2  2  6  6
3  2  4  0
4  2  7  3
'''

需要以 a 列为分组，将 a 列后所有的列（示例数据中仅演示了 b、c 两列）在组内向上位移一位。再将这个位移后的数据增加到数据右侧。

思路

先对数据用 groupby() 按 a 分组，在分组对象的 apply() 方法中编写一个函数操作位移，这样就形成了新的增加数据。

最后，用原 DataFrame 的 join() 方法合并以上数据到右侧。

代码

分组，生成新的数据：

(
    df.groupby('a')
    .apply(lambda x: x.shift(-1), include_groups=False)
)
'''
       b    c
a
1 0  3.0  4.0
  1  NaN  NaN
2 2  4.0  0.0
  3  7.0  3.0
  4  NaN  NaN
'''

上述代码的 lambda 函数中，我们用 include_groups 参数排除了 a 列，然后用 shift(-1) 向上移动一位。

接下来，我们要修改这个新数据的列标签（列名），让它不与原数据冲突，同时将这个数据赋值给一个新的变量（标识符）：

right = (
    df.groupby('a')
    .apply(lambda x: x.shift(-1), include_groups=False)
    # .set_axis(['b1', 'c1'], axis=1)
    .rename(columns=lambda x: x+'1')
    .droplevel(level=0) # 删除不需要的索引
)
right
'''
    b1   c1
0  3.0  4.0
1  NaN  NaN
2  4.0  0.0
3  7.0  3.0
4  NaN  NaN
'''

修改标签时，我们用了到 rename()，它可以写一个方法来对现有所有标签名进行处理，注释代码中还演示了硬编码修改使用的 set_axis() 方法。

最后，将新数据合并：

df.join(right)
'''
   a  b  c   b1   c1
0  1  4  2  3.0  4.0
1  1  3  4  NaN  NaN
2  2  6  6  4.0  0.0
3  2  4  0  7.0  3.0
4  2  7  3  NaN  NaN
'''

这样我们就完成了这个需求。

（完）

pandas 按分组将两列上移动增加新列

数据与需求

思路

代码

相关内容