Choose values from array based on index of argmax from another array

user_667489 0
Viewed 176

Given a dataframe with columns a1, b1 and a2, b2, I want to find the column index for the largest value out of a1, b1 and then get the value with the same relative column index out of a2, b2, as shown in the want column below:

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]})

I was able to get this far, but I'm not sure what the best approach is for the final step:

# 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))

This is a simplified version of the problem - there are actually many more columns to search through - e.g. a1-z1, a2-z2.

pandas multidimensional-array numpy python
4 Answers

For two columns, this should do:

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

For several columns:

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]

You can also replace df.iloc[...] with the corresponding column names, e.g. df.iloc[:, :numcols] by df[[a_1','b_1']]

user_42384080

We can do 

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
user_79645270

This is the solution I initially came up with:

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

I think this works, but is there a simpler way?

EDIT: It seems this only works as long as you have at most 32 columns to choose from, so other options are definitely better than this one.

user_6674890

Use, DataFrame.filter along with 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
user_128331660
Related