Как добавить новый столбец с категорией как результат группировки предыдущих?

Question

Есть DataFrame со столбцами 'floor' и 'floors_total'. Мне нужно добавить отдельный столбец с их группировкой: 'Первый', 'Последний' и 'Другой' этажи.

def floors_level(floor):
    for row in data['floors_total']:
        if floor == 1:
            return 'Первый'
        if floor == row and floor != 1:
            return 'Последний'
        return 'Другой'
        
        
data['floor_level'] = data['floor'].apply(floors_level)

Функция, которую я пробовал написать, не перебирает все значения столбца 'floors_total' как row, а останавливается на первом, подходящем под условие. С категорией 'Первый' все работает, а с остальными не получается - из 3331 строки, где 'floor' равняется 'floors_total', категории 'Последний' только 308 значений, да и те в большинстве случаев не верные (не подходят под условие 'floor' == 'floors_total').

Можете подсказать что я делаю не так?

Answer 1

Лучше воспользоваться векторизированным решением, использующим функцию np.where():

import numpy as np

def floor_cat(floor, floors_total):
    return np.where(floor==floors_total, "Последний", np.where(floor==1, "Первый", "Другой"))

Пример работы:

In [220]: df
Out[220]:
   floor  floors_total
0      1             9
1      2             9
2      3             9
3      9             9
4      1            12
5      5            12
6     12            12

In [221]: df["floor_cat"] = floor_cat(df["floor"], df["floors_total"])

In [222]: df
Out[222]:
   floor  floors_total  floor_cat
0      1             9     Первый
1      2             9     Другой
2      3             9     Другой
3      9             9  Последний
4      1            12     Первый
5      5            12     Другой
6     12            12  Последний

Answer 2

В вашей реализации существует логическая ошибка и она в том, что вы сравниваете скалярное значение floor с вектором значений data['floors_total'] в цикле.

Для справки

метод Series.apply(func) передает в функцию func значения последовательно в виде скаляров / элементов Series

Если вызвать floors_level(12) для фрейма из моего соседнего ответа, то результатом будет "Другой", вместо ожидаемого "Последний".

In [245]: floors_level(12)
Out[245]: 'Другой'

Причиной тому логика цикла, которая всегда возвращает значение на первой итерации и соответственно сравнивает floor только со значением floors_total из первой строки

Answer 3

Добил функцию без использования numpy.

def floors_level(row):
        if row['floor'] == 1:
            return 'Первый'
        if row['floor'] == row['floors_total']:
            return 'Последний'
        return 'Другой'
               
data['floor_level'] = data.apply(floors_level, axis = 1)

БЛОГ НА HUSL

Как добавить новый столбец с категорией как результат группировки предыдущих?

Ответы (3 шт):