根据另一个数组中 argmax 的索引从数组中选择值答案

【问题标题】：Choose values from array based on index of argmax from another array根据另一个数组中 argmax 的索引从数组中选择值
【发布时间】：2020-06-10 17:12:20
【问题描述】：

给定一个包含 a1、b1 和 a2、b2 列的数据框，我想找到 a1、b1 中最大值的列索引，然后从 a2、b2 中获取具有相同相对列索引的值，如显示在下面的want 列中：

import pandas as pd
import numpy as np

# Sample data
df= pd.DataFrame({'a_1':[1,2,3], 'b_1': [2,1,3], 'a_2': [3,4,7], 'b_2':[5,6,8], 'want':[5, 4, 7]})

我能够做到这一点，但我不确定最后一步的最佳方法是什么：

# Get the argmax for a1, b1
df['c'] = df[['a_1', 'b_1']].idxmax(axis=1)

# Get the column index of the argmax
df['d'] = df['c'].apply(lambda x: ['a_1', 'b_1'].index(x))

这是问题的简化版本 - 实际上还有更多列可供搜索 - 例如a1-z1, a2-z2.

【问题讨论】：

标签： python pandas numpy multidimensional-array

【解决方案1】：

对于两列，应该这样做：

df['e'] = np.where(df['a_1']>=df['b_1'], df['a_2'], df['b_2'])

对于几列：

numcols = 2
idx_max = np.argmax(df.iloc[:, :numcols].values, 1)

df['e'] = df.iloc[:,numcols:2*numcols].values[np.arange(len(df)), idx_max]

您还可以将df.iloc[...] 替换为相应的列名，例如df.iloc[:, :numcols]df[[a_1','b_1']]

【讨论】：

在我正在使用的完整示例中，有很多列，例如a1-z1 和 a2-z2 - 抱歉不清楚。
您编辑的 np.argmax 与我的相似。当我撰写答案时，我没有看到您的编辑。我删除了我的:) +1

【解决方案2】：

我们可以的

s=df[['a_1','b_1']].idxmax(1).replace(['a_1','b_1'],['a_2','b_2'])
df['value']=df.lookup(s.index,s)
df
Out[23]: 
   a_1  b_1  a_2  b_2  want  value
0    1    2    3    5     5      5
1    2    1    4    6     4      4
2    3    3    7    8     7      7

【讨论】：

这是replace 和lookup 的完美组合 :) +1

【解决方案3】：

使用DataFrame.filter 和DataFrame.lookup：

cols = df.filter(regex=r'[a-zA-Z]+_1').idxmax(1).str.rstrip('1') + '2'
df['want'] = df.filter(regex=r'[a-zA-Z]+_2').lookup(df.index, cols)

# print(df)
   a_1  b_1  a_2  b_2  want
0    1    2    3    5     5
1    2    1    4    6     4
2    3    3    7    8     7

【讨论】：

【解决方案4】：

这是我最初想出的解决方案：

df['e'] = np.choose(df['d'].values, df[['a_2', 'b_2']].transpose().values)

我认为这可行，但有更简单的方法吗？

编辑：这似乎只适用于您最多有 32 列可供选择，所以其他选项肯定比这个更好。

【讨论】：