Pandas - вычислить новый столбец с переменным вводом столбца

Question

вот в чем проблема ... Представьте в качестве примера следующий кадр данных:

df = pd.DataFrame({'col1': [1, 2, 3, 4, 5], 'col2': [3, 4, 5, 6, 7],'col3': [3, 4, 5, 6, 7],'col4': [1, 2, 3, 3, 2]})

Теперь я хотел бы добавить еще один столбец "col 5", который рассчитывается следующим образом:

если значение «col4» равно 1, тогда дайте мне соответствующее значение в столбце с индексом 1 (то есть «col2» в данном случае), если «col4» равно 2, дайте мне соответствующее значение в столбце с индексом 2(т. е. «col3» в данном случае) и т. д.

Я пробовал приведенные ниже варианты и варианты, но не могу получить правильный результат

df["col5"] = df.apply(lambda x: df.iloc[x,df[df.columns[df["col4"]]]])

Любая помощьвысоко ценится!

Answer 1

Если ваш 'col4' является индикатором индекса столбца, это будет работать:

df['col5'] = df.apply(lambda x: x[df.columns[x['col4']]], axis=1)

df

#   col1  col2  col3  col4  col5
#0     1     3     3     1     3
#1     2     4     4     2     4
#2     3     5     5     3     3
#3     4     6     6     3     3
#4     5     7     7     2     7

Answer 2

Вы можете использовать необычное индексирование с NumPy и вообще избежать цикла на уровне Python:

df['col5'] = df.iloc[:, :4].values[np.arange(df.shape[0]), df['col4']]

print(df)

   col1  col2  col3  col4  col5
0     1     3     3     1     3
1     2     4     4     2     4
2     3     5     5     3     3
3     4     6     6     3     3
4     5     7     7     2     7

Вы должны увидеть значительные преимущества в производительности для больших фреймов данных:

df = pd.concat([df]*10**4, ignore_index=True)

%timeit df.apply(lambda x: x[df.columns[x['col4']]], axis=1)       # 2.36 s per loop
%timeit df.iloc[:, :4].values[np.arange(df.shape[0]), df['col4']]  # 1.01 ms per loop