python - 获取行中条件匹配的列名

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

现在，对于每个行，我必须检查哪个列包含1，然后将此列名称记录在新列中。最终的数据框将如下所示:

        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)]