基于两个Pandas DataFrame提取Top K概率对应列的实现问题

问题原因 #

原有代码运行出错的核心原因是df2的top1、top2列存储的是字符串类型的用户ID，而df1的索引是整数类型，类型不匹配导致无法正确匹配到对应数据。

修正后可运行代码 #

import pandas as pd
import numpy as np

# 构造原始数据
df1 = pd.DataFrame({
        "user1":[1001,1002,1003,1004],
        "A" :[0.003, 0.1, 0.13, 0.23],
        "B" :[0.03, 0.3, 0.22, 0.43],
        "C" :[0.5, 0.1, 0.08, 0.04],
        "D" :[0.453,0.2,0.2,0.14]
})
df1.set_index("user1", inplace=True)

df2 = pd.DataFrame({
    "user2":[1001, 1002],
    "top1":["1004","1003"],
    "top2":["1003", "1002"]
})

# 核心修正：将top列的字符串转成整数，匹配df1索引类型
df2[['top1', 'top2']] = df2[['top1', 'top2']].astype(int)

# 转置df1
df1_T = df1.T

# 生成结果
df3 = pd.DataFrame({
    row["user2"]: np.append(
        df1_T[row["top1"]].nlargest(2).index.values,
        df1_T[row["top2"]].nlargest(2).index.values
    ) for (i, row) in df2.iterrows()
}).T

# 调整列名符合要求
df3 = df3.reset_index()
df3.columns = ['user', '1', '2', '3', '4']

print(df3)

输出结果 #

user  1  2  3  4
0  1001  B  A  B  D
1  1002  B  D  B  D

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