如何在Pandas中添加统计前一个ID行数的新列n_ID@t-1

现有如下Pandas DataFrame：

ID   Date       feature
1    2020-05-01 2 
1    2020-05-01 3 
1    2020-05-01 4
2    2019-03-15 3
2    2019-03-15 2
3    2022-04-22 5
3    2022-04-22 8
3    2022-04-22 4
3    2022-04-22 2
4    2015-01-18 4
4    2015-01-18 6
4    2015-01-18 7

执行df.sort_values('Date', ascending=False, inplace=True)按日期降序排序后，得到：

ID   Date       feature
3    2022-04-22 5
3    2022-04-22 8
3    2022-04-22 4
3    2022-04-22 2
1    2020-05-01 2 
1    2020-05-01 3 
1    2020-05-01 4
2    2019-03-15 3
2    2019-03-15 2
4    2015-01-18 4
4    2015-01-18 6
4    2015-01-18 7

需要新增列n_ID@t-1，用于记录当前ID的上一个（时间更晚的）ID的总行数，最终期望输出如下：

ID   Date       feature n_ID@t-1
3    2022-04-22 5       3   # 上一个ID为1，共3行
3    2022-04-22 8       3
3    2022-04-22 4       3
3    2022-04-22 2       3
1    2020-05-01 2       2   # 上一个ID为2，共2行
1    2020-05-01 3       2 
1    2020-05-01 4       2
2    2019-03-15 3       3   # 上一个ID为4，共3行
2    2019-03-15 2       3
4    2015-01-18 4       0   # 无后续ID，填充0
4    2015-01-18 6       0
4    2015-01-18 7       0  

已知value_counts()和transform方法，但不清楚如何组合实现需求。

可以通过以下步骤实现需求：

1. 统计每个ID的总行数 #

先用value_counts()获取每个ID对应的行数，转为字典方便后续查询：

id_row_counts = df['ID'].value_counts().to_dict()
# 输出：{3: 4, 1: 3, 2: 2, 4: 3}

2. 获取排序后的唯一ID序列 #

从排序后的DataFrame中提取去重的ID列表，顺序与分组顺序一致：

sorted_unique_ids = df['ID'].unique().tolist()
# 输出：[3, 1, 2, 4]

3. 构建ID与上一个ID行数的映射 #

遍历唯一ID列表，为每个ID匹配其下一个（时间更晚）ID的行数，最后一个ID（时间最早）填充0：

prev_id_map = {}
for idx, current_id in enumerate(sorted_unique_ids):
    if idx < len(sorted_unique_ids) - 1:
        # 当前ID的上一个ID是列表中后一位的元素
        next_id = sorted_unique_ids[idx + 1]
        prev_id_map[current_id] = id_row_counts[next_id]
    else:
        prev_id_map[current_id] = 0
# 输出：{3: 3, 1: 2, 2: 3, 4: 0}

4. 映射值到新列 #

用map()方法将映射字典的值填充到新列中：

df['n_ID@t-1'] = df['ID'].map(prev_id_map)

完整可运行代码 #

import pandas as pd

# 构造原始数据
data = {
    'ID': [1,1,1,2,2,3,3,3,3,4,4,4],
    'Date': ['2020-05-01','2020-05-01','2020-05-01','2019-03-15','2019-03-15','2022-04-22','2022-04-22','2022-04-22','2022-04-22','2015-01-18','2015-01-18','2015-01-18'],
    'feature': [2,3,4,3,2,5,8,4,2,4,6,7]
}
df = pd.DataFrame(data)

# 按日期降序排序
df.sort_values('Date', ascending=False, inplace=True)

# 统计每个ID行数
id_row_counts = df['ID'].value_counts().to_dict()

# 获取排序后的唯一ID
sorted_unique_ids = df['ID'].unique().tolist()

# 构建上一个ID行数的映射
prev_id_map = {}
for idx, current_id in enumerate(sorted_unique_ids):
    if idx < len(sorted_unique_ids) - 1:
        prev_id_map[current_id] = id_row_counts[sorted_unique_ids[idx+1]]
    else:
        prev_id_map[current_id] = 0

# 添加新列
df['n_ID@t-1'] = df['ID'].map(prev_id_map)

# 打印结果
print(df)

执行上述代码后，即可得到符合要求的DataFrame。

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