如何在熊猫数据框中实现以下输出

Question

df：

category    A   B   C   D
x   0   1   0   0
y   1   0   0   0
z   1   0   0   0
l   0   0   0   1
m   0   1   0   0
n   0   0   1   0

如何获得如下所示的df

Category    Sub-category
x   B
y   A
z   A
l   D
m   B
n   C

我尝试过：

df['sector'] = df.apply(lambda x: df.columns[x.argmax()], axis = 1)

但收到TypeError ：（“此dtype不允许进行归约运算'argmax'，'发生在索引1'）

Answer 1

就做

df['sub_category'] = df[['A', 'B', 'C', 'D']].idxmax(axis=1)


    category    A   B   C   D   sub_category
0   x           0   1   0   0   B
1   y           1   0   0   0   A
2   z           1   0   0   0   A
3   l           0   0   0   1   D
4   m           0   1   0   0   B
5   n           0   0   1   0   C

当然您可以只选择所需的列

df[['category', 'sub_category']]

    category    sub_category
0   x           B
1   y           A
2   z           A
3   l           D
4   m           B
5   n           C