在同一DataFrame中高效追踪ID变更链获取最终ID的方法

针对31k行规模的ID变更DataFrame，要高效追踪每行的最终ID，推荐使用路径压缩的映射查找法，避免循环或重复合并的高开销，具体实现如下：

示例数据 #

先还原问题中的初始DataFrame：

import pandas as pd

df = pd.DataFrame({
    "Old ID": [1, 2, 3, 6],
    "New ID": [2, 3, 6, 10]
})

实现步骤 #

1. 构建ID映射字典 #

将Old ID与New ID的对应关系转为字典，实现O(1)的快速查找：

id_map = df.set_index("Old ID")["New ID"].to_dict()

2. 带路径压缩的最终ID查找函数 #

通过递归+缓存的方式，将中途节点直接映射到最终ID，后续查询无需重复遍历链条，大幅提升效率：

def get_last_id(node, id_map, cache):
    # 当前节点无后续变更，直接返回自身
    if node not in id_map:
        return node
    # 已缓存过最终ID，直接返回
    if node in cache:
        return cache[node]
    # 递归查找最终ID
    last_id = get_last_id(id_map[node], id_map, cache)
    # 路径压缩：把当前节点直接指向最终ID，后续查询一步到位
    cache[node] = last_id
    return last_id

3. 生成Last ID列 #

初始化缓存字典，对Old ID列批量应用查找函数：

cache = {}
df["Last ID"] = df["Old ID"].apply(lambda x: get_last_id(x, id_map, cache))

最终结果 #

运行后得到目标DataFrame：

Old ID  New ID  Last ID
0       1       2       10
1       2       3       10
2       3       6       10
3       6      10       10

效率说明 #

这种方法利用路径压缩优化，每个ID最多被遍历2次（首次查找+缓存写入），时间复杂度接近O(n)，完全适配31k行的大规模数据集，不会出现循环或合并操作带来的性能瓶颈。

内容的提问来源于stack exchange，提问作者Lucio Diprè