user2920100
user2920100

Reputation: 89

Pandas row value based on column value

I have a dataframe that looks something like that:

A1   A2   A3   A4   B   C   D
0     2   9     0   9   7   2
7     6   7     3   6   8   4     
3     7   4     9   2   1   1

I want to create a new column, call it E, whose values come from columns A1, A2, A3, or A4 depending on the value of column D.

For example, if D=2, then the value for E is taken from column A2. My new dataframe would look like this:

A1   A2   A3   A4   B   C   D   E
0     2   9     0   9   7   2   2
7     6   7     3   6   8   4   3  
3     7   4     9   2   1   1   3

I have tried using df.apply but that doesn't seem to work... Any idea on how that can be done?

Thanks!

Upvotes: 0

Views: 53

Answers (1)

BENY
BENY

Reputation: 323226

We just need lookup (see documentation)

df.lookup(df.index,df.columns[df.D-1])
Out[309]: array([2, 3, 3], dtype=int64)

After assign it back

df['E']=df.lookup(df.index,df.columns[df.D-1])
df
Out[311]: 
   A1  A2  A3  A4  B  C  D  E
0   0   2   9   0  9  7  2  2
1   7   6   7   3  6  8  4  3
2   3   7   4   9  2  1  1  3

Update

df.D.astype(str).map(dict(zip(df.columns.str[-1],df.columns)))
Out[316]: 
0    A2
1    A4
2    A1
Name: D, dtype: object

Then

df.lookup(df.index,df.D.astype(str).map(dict(zip(df.columns.str[-1],df.columns))))
Out[317]: array([2, 3, 3], dtype=int64)

Upvotes: 2

Related Questions