Find pandas column names with non-zero row values

Viewed 936

I have a pandas dataframe df

df = pd.DataFrame(np.array([[0.6, 0.4, 0], [0.0, 1.0, 0], [0.4, 0.6, 0]]),columns = [-5, 0, 5])
print(df)

    -5    0    5
0  0.6  0.4  0.0
1  0.0  1.0  0.0
2  0.4  0.6  0.0

I want to find name of columns with non-zero row values. The columns names of non-zero value in each row should be stored in a list. For example following is the expected outcome for three row in df

Expected outcome:

[[-5,0],[0],[-5,0]]
3 Answers

Use:

cols = df.columns.values
mask = df.gt(0.0).values
out = [cols[x].tolist() for x in mask]

Output:

>>> out
[[-5, 0], [0], [-5, 0]]

This may be not the most efficient way. It simply goes through the entire df and records the columns.

(m,n) = df.shape
res_list = []
for i in range(m):
    temp = []
    for col in df.columns:
        if df.loc[i, col] > 0:
            temp.append(col)
    res_list.append(temp)

You can use .apply() on each row, with a lambda function to select the index (equals column index) with corresponding non-zero elements. Then, aggregate all row results into a list with the .to_list(), as follows:

df.apply(lambda x: x.index[x != 0.0].to_list(), axis=1).to_list()



# Result

[[-5, 0], [0], [-5, 0]]

Alternatively, you can also use .apply() on each row, with a lambda function together with a list comprehension going through the row elements. Then, aggregate all row results into a list with the .to_list(), as follows:

df.apply(lambda x: [y for y, z in x.items() if z != 0.0], axis=1).to_list()




# Result:

[[-5, 0], [0], [-5, 0]]
Related