如何高效将分组后DataFrame的整数索引映射回原父DataFrame？

嘿，我刚好碰到过类似的大数据量分组索引映射问题，你的方法之所以内存爆掉，主要是因为把所有XY组合转成列表再去匹配的过程，会生成大量冗余数据，完全没必要。给你一个超级高效的解决方案：

用groupby.ngroup()直接生成组索引 #

Pandas的分组对象自带了ngroup()方法，专门用来给每个分组分配一个唯一的整数索引，而且这个索引和你生成的df2d的整数索引完全对应，关键是它是内部优化过的，内存占用极低，速度超快。

直接替换你原来的映射代码就行：

# 直接在原df3d上添加分组索引，无需额外创建映射结构
df3d["idx2d"] = df3d.groupby(["X", "Y"]).ngroup()

为什么这个方法更优？ #

• 内存友好：不需要像你原来的方法那样，把几十万甚至几百万个(X,Y)元组存成列表，也不需要创建额外的映射Series，所有计算都是在分组过程中直接完成的。
• 速度更快：ngroup()是Pandas底层优化的方法，避免了手动查找匹配的开销，对于2400万行的数据集，这个方法的执行时间会比你的原方法快几个数量级。
• 结果一致：默认情况下，groupby的分组顺序和你生成df2d时的顺序完全一致（都是按分组键排序后的顺序），所以ngroup()生成的整数索引和df2d的行索引一一对应。

验证结果一致性 #

如果你担心结果和原方法有差异，可以做个小验证：

# 原方法生成的索引
idx2d = pd.Series(np.arange(len(df2d)), index=df2d.index)
original_idx = idx2d.loc[df3d[["X", "Y"]].to_records(index=False).tolist()].values

# 新方法生成的索引
new_idx = df3d.groupby(["X", "Y"]).ngroup().values

# 检查是否完全一致
print(np.array_equal(original_idx, new_idx))  # 输出True

备注：内容来源于stack exchange，提问作者Mike T