根据条件识别具有行中第一个值的列

Question

我有数据框，我想在其中标识每行的列，该列具有与条件对应的第一个值。 在下面的这种情况下，我想创建一个新列，它标识每行中小于或等于 1 的第一个值，并为我提供相应的列名。

df = pd.DataFrame({'A': [1.5,2,4,0.5], 'B' : [2,1,3,0.25], 'C': [3,1,1,1], 'D': [2,2,3,1]})
df
    A    B      C   D
0   1.5  2.00   3   2
1   2.0  1.00   1   2
2   4.0  3.00   1   3
3   0.5  0.25   1   1

我可以创建一个掩码来检查条件。

temp = df<=1
temp
    A       B       C       D
0   False   False   False   False
1   False   True    True    False
2   False   False   True    False
3   True    True    True    True

然后我可以使用以下内容来确定列。

df['New_col'] = temp.idxmax(axis = 1)
df
    A    B      C   D   New_col
0   1.5  2.00   3   2   A
1   2.0  1.00   1   2   B
2   4.0  3.00   1   3   C
3   0.5  0.25   1   1   A

代码正确识别 New_col 中的列，除了第 0 行，因为第 0 行中的所有值都大于 1。如何为 New_col 中的第 0 行获取 NaN 而不是 A？

下面是所需的输出。

    A    B      C   D   New_col
0   1.5  2.00   3   2   NaN
1   2.0  1.00   1   2   B
2   4.0  3.00   1   3   C
3   0.5  0.25   1   1   A

谢谢。

Answer 1

使用any(1)检查是否有True的行，并where掩盖：

df['New_col'] = temp.idxmax(axis = 1).where(temp.any(1))

输出：

     A     B  C  D New_col
0  1.5  2.00  3  2     NaN
1  2.0  1.00  1  2       B
2  4.0  3.00  1  3       C
3  0.5  0.25  1  1       A