python - 根据行值条件从数据框中提取列名

Question

我有一个数据框

     A     B     C
 u1  0     .5    .2
 u2  .2     0     .3
 u3   .1    0     0

我需要针对每个值不为零的索引查找列名所以我需要输出

        elements 
    u1  [B,C]
    u2  [A,C]
    u3  [A]

我可以使用 df.idxmax(axis=1) 跨行查找最高值的列名但是如何找到列的所有名称。

Answer 1

您可以使用 apply使用 axis=1 按行处理并通过将值转换为 bool - 0 为 False，而不是 0 是 True:

df = df.apply(lambda x: x.index[x.astype(bool)].tolist(), 1)
print (df)
u1    [B, C]
u2    [A, C]
u3       [A]
dtype: object

如果输出应该是strings:

s = np.where(df, ['{}, '.format(x) for x in df.columns], '')
df = pd.Series([''.join(x).strip(', ') for x in s], index=df.index)
print (df)
u1    B, C
u2    A, C
u3       A
dtype: object

详细信息:

print (s)
[['' 'B, ' 'C, ']
 ['A, ' '' 'C, ']
 ['A, ' '' '']]