dataframeの列同士を比較して条件を満たす列ラベルを取り出す


質問しようと思ってあちこちいじってたら上手く行ったので、備忘録がてら登録。

主題

表1のデータから表2を作りたい。

表1 生データ

総額 食費 光熱費 水道費
2000 2000 0 0
1000 0 1000 0
1500 0 0 1500

表2 作りたいデータ

総額 食費 光熱費 水道費 分類
2000 2000 0 0 食費
1000 0 1000 0 光熱費
1500 0 0 1500 水道費

食費・光熱費・水道費はどれか一つに値が入るようになっている
食費1000+水道費1000=総額2000、というようなことはない。

自分で試したこと

import pandas as pd
df = pd.DataFrame(
    data=[{'総額': 2000, '食費': 2000,'光熱費': 0,'水道費': 0,},
          {'総額': 1000, '食費': 0,'光熱費': 1000,'水道費': 0,},
          {'総額': 1500, '食費': 0,'光熱費': 0,'水道費': 1500,}])

df_temp= df.drop("総額",axis=1)
df_temp["分類"]=df_temp.idxmax(axis=1)
df=df.merge(df_temp)

今回は3列の中の最大値が欲しかったので、idxmaxで一発で列ラベルを拾えた。
idxmax、idxminが使えるならこれが良さそう