Блог

У меня есть dataframe, где я хочу идентифицировать столбец для каждой строки, который имеет первое значение, соответствующее условию. В этом случае ниже я хочу создать новый столбец, который идентифицирует первое значение в каждой строке, которое меньше или равно 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.
Как я могу получить NaN вместо A для строки 0 в New_col?

Ниже приведен желаемый результат.

     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
  

Спасибо.

Используется any(1) для проверки наличия a 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