如何基于DataFrame的ID列值批量生成对应索引列表？

示例数据构建 #

import pandas as pd

data={'ID':[1,1,2,3,6,4,2,6], 'Number': [10,20,5,6,100,90,40,5]}
df=pd.DataFrame(data)

手动方法的局限 #

手动指定单个ID提取索引的方式（如下代码）只能处理已知的ID值，无法应对ID值数量未知的场景：

idx_list=df.index[df.ID == 1].tolist()

自动分组解决方案 #

可以利用groupby按ID分组，直接通过indices属性获取每组对应的索引数组，最终得到以ID为键、索引数组为值的字典：

id_to_indices = df.groupby('ID').indices

执行后结果为：

{1: array([0, 1]), 2: array([2, 6]), 3: array([3]), 4: array([5]), 6: array([4, 7])}

如果需要将索引转为普通列表格式，可通过字典推导式进一步处理：

id_to_indices = {k: v.tolist() for k, v in df.groupby('ID').indices.items()}

此时结果为：

{1: [0, 1], 2: [2, 6], 3: [3], 4: [5], 6: [4, 7]}

这种方式无需提前知晓ID的数量，可自动为每个不同的ID值生成对应的索引列表。

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