如何用更优雅的Pandas方法获取用户当前与上一版套餐信息？

需求 #

获取每个用户的当前套餐与上一版套餐信息。

给定数据集 #

import pandas as pd

d = {'user_id': [1, 1, 1, 3, 2, 3], 'tariff_id': [1, 2, 3, 4, 3, 1], 
'date_change':['2022-01-01', '2023-01-01', '2023-02-10', '2023-03-01', '2022-05-07', '2023-06-08']}
change_tariff = pd.DataFrame(data=d)

原有实现方式 #

change_tariff_sort = change_tariff.sort_values(by='date_change')

answ = pd.DataFrame()
for i in pd.unique(change_tariff_sort.user_id): 
    answ = answ.append(change_tariff_sort[change_tariff_sort.user_id==i].tail(2))

更优雅的Pandas解决方案 #

方式1：保留每个用户最后两条记录（分两行展示） #

利用groupby结合tail，无需循环即可完成，同时避免了已过时的append方法：

result = (change_tariff
          .sort_values(by=['user_id', 'date_change'])
          .groupby('user_id')
          .tail(2))

先按用户和日期排序，再分组取每组最后2条，直接得到每个用户的当前与上一版套餐记录。

方式2：将当前与上一版套餐合并到同一行 #

如果需要更直观的对比展示，可以用shift生成上一版套餐的字段，再保留每个用户的最新记录：

# 先按用户和日期排序
sorted_df = change_tariff.sort_values(by=['user_id', 'date_change'])

# 分组后生成上一版套餐的ID与变更日期
sorted_df['prev_tariff_id'] = sorted_df.groupby('user_id')['tariff_id'].shift(1)
sorted_df['prev_date_change'] = sorted_df.groupby('user_id')['date_change'].shift(1)

# 只保留每个用户的最新记录，一行包含当前与上一版信息
final_result = sorted_df.groupby('user_id').tail(1)

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