熊猫-使用可变列输入计算新列

时间:2018-11-29 12:55:03

标签: python pandas

问题出在这里...想象一下以下数据框为例:

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

非常感谢您的帮助!

2 个答案:

答案 0 :(得分:2)

如果您的'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

答案 1 :(得分:1)

您可以对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