如何使用pandas df.shift实现循环移位，用首尾元素填补空缺？

问题描述 #

现有如下5行2列的Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    'A': [10, 20, 30, 40, 50],
    'B': [0, 5, 10, 15, 20]
})

原始数据展示：

A   B
0  10   0
1  20   5
2  30  10
3  40  15
4  50  20

需要实现两种循环移位效果：

• 对A列执行向下2行的移位，顶部空缺用原本会被移出底部的2个元素填充
• 对B列执行向上2行的移位，底部空缺用原本会被移出顶部的2个元素填充

使用普通的df["A"].shift(2)和df["B"].shift(-2)会产生NaN值，不符合需求，预期的最终结果为：

A   B
0  40  10
1  50  15
2  10  20
3  20   0
4  30   5

解决方案 #

可以借助numpy.roll()函数实现循环移位，该函数会将数组元素进行循环滚动，自动用移出的元素填充空缺位置：

import pandas as pd
import numpy as np

# 创建原始DataFrame
df = pd.DataFrame({
    'A': [10, 20, 30, 40, 50],
    'B': [0, 5, 10, 15, 20]
})

# A列向下循环移2行：将末尾2个元素移到开头
df['A'] = np.roll(df['A'], shift=2)
# B列向上循环移2行：将开头2个元素移到末尾
df['B'] = np.roll(df['B'], shift=-2)

print(df)

代码说明 #

• np.roll(arr, shift=n)：当n为正数时，将数组arr的最后n个元素移动到数组开头，对应DataFrame列的向下循环移位；
• 当n为负数时，将数组arr的开头|n|个元素移动到数组末尾，对应DataFrame列的向上循环移位；
• 该方法不会产生NaN值，完美实现循环填充的移位效果。

内容的提问来源于stack exchange，提问作者hbstha123