如何根据val列最大值获取对应id列值？Pandas数据处理需求

问题背景 #

现有如下Pandas DataFrame：

import pandas as pd
data = {"id_1":["a","b","c"],
        "id_2":["q","w","e"],
        "val_1":[1,2,3],
        "val_2":[2,0,0]}
df = pd.DataFrame(data)

初始数据输出：

id_1 id_2  val_1  val_2
0    a    q      1      2
1    b    w      2      0
2    c    e      3      0

需求：比较val_1和val_2列的值（二者永不相等），若val_1大于val_2，新列max_id取id_1的值；否则取id_2的值，最终目标结果：

id_1 id_2  val_1  val_2 max_id
0    a    q      1      2      q
1    b    w      2      0      b
2    c    e      3      0      c

用户已实现获取最大值及对应id列名，但不清楚如何生成最终的max_id列。

解决方案 #

方法1：基于已实现的列名映射取值 #

你已经拿到了每行对应的目标id列名，直接用df.lookup按行取值即可：

index_map = {0:"id_1", 1:"id_2"}
# 获取每行对应的id列名
target_cols = df.apply(lambda x: [x.val_1, x.val_2].index(max([x.val_1, x.val_2])), axis=1).map(index_map)
# 生成max_id列
df['max_id'] = df.lookup(df.index, target_cols)

方法2：更高效的np.where方法（推荐） #

无需绕弯子找列名，直接通过条件判断取值，效率远高于逐行apply：

import numpy as np
df['max_id'] = np.where(df['val_1'] > df['val_2'], df['id_1'], df['id_2'])

方法3：使用Pandas原生where方法 #

逻辑与np.where一致，用Pandas内置方法实现：

df['max_id'] = df['id_1'].where(df['val_1'] > df['val_2'], df['id_2'])

以上三种方法均可得到目标结果，其中方法2和3避免了逐行遍历，在数据量较大时性能更优。

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