pandas 按月数滞后多索引不规则时间序列数据

Question

我有以下熊猫数据框

df = pd.DataFrame(data = {
    'item': ['red','red','red','blue','blue'],
    'dt': pd.to_datetime(['2018-01-31', '2018-02-28', '2018-03-31', '2018-01-31', '2018-03-31']),
    's': [3.2, 4.8, 5.1, 5.3, 5.8],
    'r': [1,2,3,4,5],
    't': [7,8,9,10,11],
})

看起来像

    item    dt          s   r   t
0   red     2018-01-31  3.2 1   7
1   red     2018-02-28  4.8 2   8
2   red     2018-03-31  5.1 3   9
3   blue    2018-01-31  5.3 4   10
4   blue    2018-03-31  5.8 5   11

请注意，时间点不规则：“蓝色”缺少 2 月数据。所有日期均为有效的月末日期。

我想添加一个列，它是“两个月前的 s 值”，理想情况下是这样的

df['s_lag2m'] = df.set_index(['item','dt'])['s'].shift(2, 'M')

我会得到

    item    dt          s   r   t   s_lag2m
0   red     2018-01-31  3.2 1   7   NaN
1   red     2018-02-28  4.8 2   8   NaN
2   red     2018-03-31  5.1 3   9   3.2
3   blue    2018-01-31  5.3 4   10  NaN
4   blue    2018-03-31  5.8 5   11  5.3

但这不起作用；它抛出NotImplementedError: Not supported for type MultiIndex。

我该怎么做？

Answer 1

我们可以在set_index 之后只用 dt 做reindex

df['New']=df.set_index(['dt']).groupby('item')['s'].shift(2, 'M').\
             reindex(pd.MultiIndex.from_frame(df[['item','dt']])).values
df
   item         dt    s  r   t  New
0   red 2018-01-31  3.2  1   7  NaN
1   red 2018-02-28  4.8  2   8  NaN
2   red 2018-03-31  5.1  3   9  3.2
3  blue 2018-01-31  5.3  4  10  NaN
4  blue 2018-03-31  5.8  5  11  5.3