Получение имени столбца, где условие соответствует строке

Question

У меня есть pandas кадр данных, который выглядит следующим образом:

        A     B     C     D     E     F     G     H     I
1       0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
2       1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0
3       0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0

Теперь для каждого row я должен проверить, какое column содержит 1, а затем записать имя этого столбца в новый столбец. Финал dataframe будет выглядеть так:

        A     B     C     D     E     F     G     H     I     IsTrue
1       0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   B
2       1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   A
3       0.0   1.0   0.0   0.0   0.0   0.0   0.0   0.0   0.0   B

Есть ли какой-нибудь более быстрый и pythonic способ сделать это?

Answer 1

Вот один из способов использования DataFrame.dot:

df['isTrue'] = df.astype(bool).dot(df.columns)

    A    B    C    D    E    F    G    H    I    isTrue
1  0.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0      B
2  1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0      A
3  0.0  1.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0      B

---

Для еще лучшей производительности вы можете использовать:

df['isTrue'] = df.columns[df.to_numpy().argmax(1)]

Answer 2

Что вы описали, это определение idxmax

>>> df.idxmax(1)

1    B
2    A
3    B
dtype: object