如何获取关联不同日期相同id的DataFrame索引字典？

示例数据 #

示例DataFrame结构： #

| id| date
--------------
0   | a | d1
1   | b | d1
2   | a | d2
3   | c | d2
4   | b | d3
5   | a | d3

代码创建DataFrame： #

import pandas as pd

df = pd.DataFrame({'id': ['a', 'b', 'a', 'c', 'b', 'a'], 'date': ['d1', 'd1', 'd2', 'd2', 'd3', 'd3']})

需求说明 #

需要生成一个链接字典，例如 d = {0: 2, 2: 5, 1: 4}，其中键值对含义：

• 0:2：d1中的a（行索引0）关联到d2中的a（行索引2）
• 2:5：d2中的a（行索引2）关联到d3中的a（行索引5）
• 1:4：d1中的b（行索引1）关联到d3中的b（行索引4）

实现方法 #

推荐两种简洁清晰的实现方式：

方式一：分组遍历索引列表（直观易懂） #

通过groupby按id分组，获取每个id对应的所有行索引，再遍历每组索引列表，将相邻索引配对：

# 按id分组，获取每个id对应的行索引列表
grouped_indices = df.groupby('id').indices

link_dict = {}
# 遍历每个id的索引列表，配对相邻索引
for idx_list in grouped_indices.values():
    # 只处理有多个记录的id
    if len(idx_list) >= 2:
        for i in range(len(idx_list)-1):
            link_dict[idx_list[i]] = idx_list[i+1]

print(link_dict)
# 输出：{0: 2, 2: 5, 1: 4}

方式二：链式操作（简洁高效） #

利用groupby+shift的组合，直接生成映射关系：

link_dict = (
    df.groupby('id')
      .shift(-1)
      .dropna()
      .index.to_series()
      .map(df.groupby('id').shift(-1).index)
      .dropna()
      .to_dict()
)

print(link_dict)
# 输出：{0: 2, 2: 5, 1: 4}

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