Распределение столбцов датафрейма по группам числовых и категориальных типов

Рассмотрим, как быстро разделить колонки датафрейма на категориальные и числовые группы. Это является важным этапом перед преобразованием типов и формы подачи данных в алгоритм машинного обучения, так как подходы к их обработке зачастую разнятся. Загрузим тестовый датасет для демонстрации:

from sklearn.datasets import load_diabetes

df,_ = load_diabetes(return_X_y=True, as_frame=True)
display(df.head())

Базовую часть работы будет выполнять метод датафрейма, который возвращает серию с информацией о количестве уникальных значений в каждом столбце (по умолчанию)/строке:

df.nunique()

Теперь установим уровень количества уникальных значений, который будет регулировать, считаем ли мы рассматриваемую колонку категориальной (если меньше или равно) или числовой. Затем на основании этой условной границы разделим столбцы по группам:

cat_level = 30
cols_ser = df.nunique()

cat_cols = cols_ser[cols_ser<=cat_level].index.tolist()
num_cols = cols_ser[cols_ser>cat_level].index.tolist()
display(cat_cols)
display(num_cols)

В коде предполагается, что на этапе формирования cols_ser в df присутствуют только столбцы, относящиеся к указанным категориям. Если есть лишние, то удалите их на этом этапе, используя метод drop.