我有一个数据框
col0 col1 col2 col3 col4
0 1 3 6 6 0
1 0 2 8 7 3
2 0 0 4 3 4
3 4 2 2 0 4
逻辑是如果 col1 不为零,则返回 col1。如果 col 1 为零,则返回 col2(非零)。如果 col 2 为零,则返回 col3。我们不需要为 col4 做任何事情
我的代码如下所示,但它只返回 col1
def test(df):
if df['col1'].iloc[0] > 0:
return df['col1']
elif df['col1'].iloc[0] == 0 & df['col2'].iloc[0] > 0:
return df['col2']
elif df['col2'].iloc[0] == 0 & df['col3'].iloc[0] > 0:
return df['col3']
else:
return 0
test(new)
我尝试了 .any() 和 .all(),它也不起作用。另外,有没有办法让这段代码更高效?
答案 0 :(得分:1)
@ALollz 想法的变体,因为 lookup 在 pandas 1.2.0 中已弃用:
indices = np.argmax(df.ne(0).values, axis=1)
print(df.values[np.arange(len(df)), indices])
输出
[1 2 4 4]
更新
要排除最后一列并返回 0,请改为执行此操作:
indices = np.argmax(df.ne(0).iloc[:, :-1].values, axis=1)
result = np.where(df.ne(0).iloc[:, :-1].any(1), df.values[np.arange(len(df)), indices], 0)
print(result)