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

获取条件行连续匹配的列名

2 个答案: