使用Pandas获取行中第一个非零值所在的列名

Question

我有一个巨大的数据框，但只分享以下的样本。它是一个带有如下所示的样本标题列名的CSV文件。

sample.csv
cnum,sup1,sup2,sup3,sup4
285414459,1,0,1,1
445633709,1,0,0,0
556714736,0,0,1,0
1089852074,0,1,0,1

一个 cnum 在所有 sup* 列中可以有0个或1个设置。我想选择并打印出遇到第一个 1 的列名，之后出现的所有其他1都应该被忽略，输出中不应该打印任何列名。

expected output:
cnum,supcol
285414459,sup1
445633709,sup1
556714736,sup3
1089852074,sup2

目前我尝试了这段代码：

import pandas as pd
df=pd.read_csv('sample.csv')
df_union=pd.DataFrame(columns=['cnum','supcol'])
for col in df.columns: 
    df1=df.filter(['cnum']).loc[df[col] == 1]
    df1['supcol']=col
    df_union=df_union.append(df1)
print(df_union)

然而它打印了所有列名，其中该列名的值为1。我只想要第一个。请帮忙。

- Pallavi Chaporkar

Answer 1

看起来您可以在此处使用 idxmax：

df.set_index('cnum').idxmax(axis=1).reset_index(drop=True)

0    sup1
1    sup1
2    sup3
3    sup2
dtype: object

df['output'] = df.set_index('cnum').idxmax(axis=1).reset_index(drop=True) 
# Slightly faster,
# df['output'] = df.set_index('cnum').idxmax(axis=1).to_numpy() 

df
         cnum  sup1  sup2  sup3  sup4 output
0   285414459     1     0     1     1   sup1
1   445633709     1     0     0     0   sup1
2   556714736     0     0     1     0   sup3
3  1089852074     0     1     0     1   sup2

使用 dot 的另一个选项（将为您提供所有非零列）：

d = df.set_index('cnum') 
d.dot(d.columns + ',').str.rstrip(',').reset_index(drop=True)

0    sup1,sup3,sup4
1              sup1
2              sup3
3         sup2,sup4
dtype: object

或者，

(d.dot(d.columns + ',')
  .str.rstrip(',')
  .str.split(',', 1).str[0] 
  .reset_index(drop=True))

0    sup1
1    sup1
2    sup3
3    sup2
dtype: object

- cs95