Значение строки панд на основе значения столбца

Question

У меня есть датафрейм, который выглядит примерно так:

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

Я хочу создать новый столбец, назовите его E, значения которого берутся из столбцов A1, A2, A3 или A4 в зависимости от значения столбца D.

Например, если D = 2, то значение для E берется из столбца A2. Мой новый фрейм данных будет выглядеть так:

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

Я пытался использовать df.apply, но, похоже, это не работает ... Есть идеи, как это можно сделать?

Спасибо!

Answer 1

Нам просто нужно lookup (см. документацию )

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

После присвоения обратно

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

Обновление

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

Тогда

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)