I have a dataframe
JavaScript
x
6
1
col0 col1 col2 col3 col4
2
0 1 3 6 6 0
3
1 0 2 8 7 3
4
2 0 0 4 3 4
5
3 4 2 2 0 4
6
The logic is if col1 is not zero, return col1. If col 1 is zero, return col2 (non-zero). If col 2 is zero, return col3. We don’t need to do anything for col4
My code looks like below but it only returns col1
JavaScript
1
11
11
1
def test(df):
2
if df['col1'].iloc[0] > 0:
3
return df['col1']
4
elif df['col1'].iloc[0] == 0 & df['col2'].iloc[0] > 0:
5
return df['col2']
6
elif df['col2'].iloc[0] == 0 & df['col3'].iloc[0] > 0:
7
return df['col3']
8
else:
9
return 0
10
test(new)
11
I tried .any() and .all(), it doesnt work either. Also, is there anyway to make this piece of code more efficient?
Advertisement
Answer
A variation on @ALollz idea, since lookup is deprecated on pandas 1.2.0:
JavaScript
1
3
1
indices = np.argmax(df.ne(0).values, axis=1)
2
print(df.values[np.arange(len(df)), indices])
3
Output
JavaScript
1
2
1
[1 2 4 4]
2
UPDATE
For excluding the last column, and return 0, do this instead:
JavaScript
1
4
1
indices = np.argmax(df.ne(0).iloc[:, :-1].values, axis=1)
2
result = np.where(df.ne(0).iloc[:, :-1].any(1), df.values[np.arange(len(df)), indices], 0)
3
print(result)
4