如何根据DataFrame的y列将pairs列列表元素替换为对应id值？

需求说明 #

保持id列不变，以y列为映射键，将pairs列中的每个元素替换为对应y值所在行的id。

输入DataFrame #

import pandas as pd

df_input = pd.DataFrame({
    'id': [111, 222, 333, 444],
    'y': [0, 1, 3, 4],
    'pairs': [[0], [1, 4], [0, 3], [4]]
})

期望输出 #

df_output = pd.DataFrame({
    'id': [111, 222, 333, 444],
    'y': [0, 1, 3, 4],
    'pairs': [[111], [222, 444], [111, 333], [444]]
})

尝试的错误代码 #

def transform_df(df):
    id_to_list_map = df.set_index('y')['pairs'].to_dict()

    def modify_list(lst, id_value):
        return [id_value if i in id_to_list_map else i for i in lst]

    df2 = df.copy()
    df2['list'] = df2.apply(lambda row: modify_list(row['pairs'], row['id']), axis=1)

    return df2

错误代码输出 #

pd.DataFrame({
    'id': [111, 222, 333, 444],
    'y': [0, 1, 3, 4],
    'pairs': [[111], [444, 444], [333, 333], [444]]
})

问题分析 #

代码逻辑存在两处核心错误：

1. 构建的id_to_list_map映射关系完全错误，应该是**y值对应id**，而非y值对应pairs列表；
2. modify_list函数用当前行的id替换所有存在于映射中的元素，导致不管pairs里的元素对应哪个y值，都被替换成当前行的id，这是错误结果的直接原因。

正确解法 #

方法1：构建y到id的映射字典，批量替换 #

def transform_df(df):
    # 构建正确的y值到id的映射字典
    y_to_id = df.set_index('y')['id'].to_dict()
    
    # 遍历每个pairs列表，替换元素为对应id
    df2 = df.copy()
    df2['pairs'] = df2['pairs'].apply(lambda lst: [y_to_id[num] for num in lst])
    
    return df2

方法2：拆分逻辑的写法（与方法1逻辑一致） #

def transform_df(df):
    y_to_id = df.set_index('y')['id'].to_dict()
    
    def replace_pairs(lst):
        return [y_to_id[val] for val in lst]
    
    df2 = df.copy()
    df2['pairs'] = df2['pairs'].apply(replace_pairs)
    return df2

验证结果 #

调用函数后输出与期望完全一致：

print(transform_df(df_input))
# 输出：
    id  y       pairs
0  111  0       [111]
1  222  1  [222, 444]
2  333  3  [111, 333]
3  444  4       [444]

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