5 мощных способов добавления столбцов с условиями в pandas

Для кого эта статья:

• Аналитики данных и специалисты по обработке данных
• Студенты и начинающие разработчики, изучающие Python и библиотеки для работы с данными

• Профессионалы, стремящиеся повысить производительность своих аналитических процессов и освоить эффективные техники работы с pandas

  Манипулирование данными — это хлеб и масло каждого аналитика. При работе с pandas часто возникает необходимость создавать новые столбцы, значения которых зависят от нескольких условий. Будь то категоризация клиентов по сложным параметрам или трансформация числовых данных в бизнес-метрики — владение техниками условного добавления столбцов критически важно для эффективной обработки данных. В этой статье я разберу 5 мощных способов добавления столбцов с множественными условиями в pandas, от классического numpy.where до сложных функциональных подходов. Выбор правильного метода может ускорить ваш код в десятки раз! 🔥

Хотите стать мастером обработки данных в Python и легко решать сложные задачи с pandas? На курсе Обучение Python-разработке от Skypro вы не только освоите базовые и продвинутые методы работы с данными, но и научитесь создавать элегантные решения для бизнес-задач. Наши студенты уже применяют условные преобразования в реальных проектах, значительно ускоряя аналитические процессы. Присоединяйтесь!

Создание столбцов в Pandas с множественными условиями

Работая с данными, мы часто сталкиваемся с необходимостью классифицировать информацию по нескольким условиям. Например, вам может потребоваться создать столбец "Ценовой сегмент", который определяется не только ценой товара, но и его категорией, или сегментировать клиентов по комбинации возраста, дохода и частоты покупок.

Pandas предлагает несколько элегантных способов добавления столбцов на основе множественных условий. Давайте рассмотрим их на примере простого датафрейма с данными о продажах:

import pandas as pd
import numpy as np

# Создаем пример DataFrame
data = {
'Product': ['A', 'B', 'C', 'A', 'B'],
'Price': [100, 200, 50, 300, 150],
'Category': ['Electronics', 'Furniture', 'Books', 'Electronics', 'Furniture'],
'Sales': [10, 5, 20, 8, 15]
}

df = pd.DataFrame(data)

Допустим, нам нужно добавить столбец 'PriceCategory', который будет иметь значение:

• 'Premium' — если цена > 200 и категория 'Electronics'
• 'High' — если цена > 100 но <= 200, или если категория 'Furniture' и цена > 100
• 'Low' — во всех остальных случаях

Существует несколько способов решить эту задачу, каждый со своими преимуществами и недостатками. Важно выбрать подходящий метод в зависимости от сложности условий, размера данных и требований к читаемости кода.

Александр Петров, Lead Data Analyst

В прошлом году я работал над проектом оптимизации логистики крупного интернет-магазина. Нам требовалось срочно сегментировать десятки тысяч SKU по сложной системе приоритетов доставки. Критерии включали сразу несколько параметров: маржинальность товара, вес, габариты, среднюю частоту заказов и стоимость.

Сначала я использовал каскад условных операторов через метод .loc, но время выполнения было неприемлемо большим — около 4 минут на полном датасете. Когда я переписал логику с использованием np.select и векторизованных операций, время сократилось до 6 секунд! Это позволило интегрировать обработку в микросервис, работающий в режиме реального времени.

Правильно выбранный метод добавления столбцов с множественными условиями может кардинально повлиять на производительность вашего приложения. Важно понимать все доступные варианты и их применимость в разных ситуациях.

Теперь давайте рассмотрим каждый метод подробнее, начиная с наиболее часто используемых функций из библиотеки NumPy.

Метод np.where() и np.select() для добавления столбцов

Функции numpy.where() и numpy.select() — это мощные инструменты для добавления столбцов на основе условий. Они работают непосредственно с массивами, что делает их очень эффективными с точки зрения производительности. 💪

Использование np.where()

Функция np.where() работает по принципу тернарного оператора: если условие верно, возвращается одно значение, иначе — другое:

# Добавление столбца с одним условием
df['is_expensive'] = np.where(df['Price'] > 150, 'Yes', 'No')

Для множественных условий np.where() можно вкладывать друг в друга:

# Вложенное использование np.where для множественных условий
df['PriceCategory'] = np.where(
(df['Price'] > 200) & (df['Category'] == 'Electronics'), 
'Premium',
np.where(
((df['Price'] > 100) & (df['Price'] <= 200)) | 
((df['Category'] == 'Furniture') & (df['Price'] > 100)),
'High',
'Low'
)
)

Использование np.select()

Если условий много, код с вложенными np.where() становится сложным для чтения. В таких случаях лучше использовать np.select(), который принимает списки условий и соответствующих значений:

# Создание списков условий и значений
conditions = [
(df['Price'] > 200) & (df['Category'] == 'Electronics'),
((df['Price'] > 100) & (df['Price'] <= 200)) | 
((df['Category'] == 'Furniture') & (df['Price'] > 100))
]
choices = ['Premium', 'High']
default = 'Low'

# Применение np.select
df['PriceCategory'] = np.select(conditions, choices, default)

Преимущества np.select() становятся очевидны, когда количество условий превышает два-три. Код становится более структурированным и легче поддерживаемым.

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

Важно помнить, что условия в np.select() проверяются по порядку, и как только находится первое истинное условие, соответствующее значение выбирается, а остальные условия игнорируются. Это может быть как преимуществом, так и источником ошибок, если порядок условий не соответствует вашей логике.

Использование loc[] и query() при работе с условиями

Методы .loc[] и .query() представляют собой альтернативный подход к добавлению столбцов с условиями, который многие разработчики находят более интуитивным и читаемым. Эти методы особенно полезны, когда вам нужно не только создать новый столбец, но и модифицировать существующие данные на основе условий. 🧩

Использование .loc[]

Метод .loc[] позволяет нам выбирать строки и столбцы по их меткам, но также может использоваться для изменения значений в датафрейме на основе условных выражений:

# Создаем новый столбец со значением по умолчанию
df['PriceCategory'] = 'Low'

# Применяем условия с помощью .loc[]
df.loc[(df['Price'] > 200) & (df['Category'] == 'Electronics'), 'PriceCategory'] = 'Premium'
df.loc[
(((df['Price'] > 100) & (df['Price'] <= 200)) | 
((df['Category'] == 'Furniture') & (df['Price'] > 100))),
'PriceCategory'
] = 'High'

Преимущество этого подхода в том, что он позволяет применять независимые условия в любом порядке, что делает код более гибким и модульным. Также .loc[] можно использовать для обновления нескольких столбцов одновременно:

# Обновление нескольких столбцов по одному условию
df.loc[df['Price'] > 200, ['PriceCategory', 'MarketingSegment']] = ['Premium', 'Target']

Использование .query()

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

# Создаем новый столбец со значением по умолчанию
df['PriceCategory'] = 'Low'

# Применяем условия с помощью .query()
premium_filter = df.query('Price > 200 and Category == "Electronics"').index
high_filter = df.query('(Price > 100 and Price <= 200) or (Category == "Furniture" and Price > 100)').index

df.loc[premium_filter, 'PriceCategory'] = 'Premium'
df.loc[high_filter, 'PriceCategory'] = 'High'

Комбинируя .query() для фильтрации и .loc[] для присваивания, мы получаем очень выразительный код, который легко поддерживать и расширять. Однако стоит помнить, что этот подход может быть медленнее векторизованных операций numpy для больших датафреймов.

Мария Соколова, Data Science Engineer

Когда я начинала работать с pandas, я предпочитала метод .loc[] для всех операций с условиями – он казался мне наиболее интуитивным. Однажды я столкнулась с задачей обработки исторических биржевых данных, где нужно было классифицировать паттерны движения цен по десяткам технических индикаторов.

Датафрейм содержал более 3 миллионов строк, и мой первоначальный подход с цепочкой операций .loc[] работал почти 15 минут. Коллега предложил переписать решение с использованием np.select(), и производительность улучшилась драматически – до 20 секунд!

Однако был интересный момент. Когда позже нам потребовалось добавить ещё один столбец с более сложной логикой, включающей вычисление на основе предыдущих значений, именно комбинирование обоих методов дало наилучший результат: np.select() для базовых условий и точечное применение .loc[] для специальных случаев.

Функциональный подход: apply() и lambda для новых столбцов

Функциональный подход с использованием .apply() и lambda-функций предлагает максимальную гибкость при создании новых столбцов на основе сложных условий. Это особенно полезно, когда логика слишком сложна для выражения через векторизованные операции или когда требуются дополнительные вычисления. 🔄

Использование .apply() с lambda-функциями

Метод .apply() позволяет применить функцию к каждой строке или столбцу датафрейма. В сочетании с анонимными lambda-функциями это дает нам компактный способ выразить сложную логику:

# Добавление столбца с помощью apply и lambda
df['PriceCategory'] = df.apply(
lambda row: 'Premium' if row['Price'] > 200 and row['Category'] == 'Electronics'
else 'High' if ((row['Price'] > 100 and row['Price'] <= 200) or 
(row['Category'] == 'Furniture' and row['Price'] > 100))
else 'Low',
axis=1
)

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

Использование именованных функций

Для еще более сложной логики лучше определить отдельную функцию вместо lambda:

def determine_price_category(row):
"""Определяет ценовую категорию на основе сложных правил"""
if row['Price'] > 200 and row['Category'] == 'Electronics':
return 'Premium'
elif ((row['Price'] > 100 and row['Price'] <= 200) or 
(row['Category'] == 'Furniture' and row['Price'] > 100)):
# Здесь можно добавить дополнительную логику, вычисления и т.д.
return 'High'
else:
# Можно использовать более сложную логику для значения по умолчанию
if row['Sales'] > 15: # Дополнительное условие
return 'Popular-Low'
return 'Low'

# Применение функции ко всему DataFrame
df['PriceCategory'] = df.apply(determine_price_category, axis=1)

Это решение более читаемо и поддерживаемо для сложной логики, особенно когда присвоение категорий требует нетривиальных вычислений или проверок с несколькими уровнями вложенности.

Метод map() для простых преобразований

Для более простых случаев, когда условие зависит от одного столбца, можно использовать метод map():

# Создаем словарь соответствия категорий
category_mapping = {
'Electronics': 'Tech',
'Furniture': 'Home',
'Books': 'Media'
}

# Применяем маппинг
df['CategoryGroup'] = df['Category'].map(category_mapping)

Этот подход особенно удобен для преобразования категориальных переменных и менее гибок для сложных условных выражений.

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

Сравнение производительности методов добавления столбцов

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

Для объективного сравнения я создам большой тестовый датафрейм и измерю время выполнения каждого метода:

import pandas as pd
import numpy as np
import time

# Создаем большой тестовый датафрейм
n_rows = 1_000_000
test_df = pd.DataFrame({
'A': np.random.randint(0, 100, n_rows),
'B': np.random.choice(['X', 'Y', 'Z'], n_rows),
'C': np.random.randn(n_rows)
})

# Функция для измерения времени выполнения
def time_execution(func, *args, **kwargs):
start = time.time()
result = func(*args, **kwargs)
end = time.time()
return result, end – start

Теперь сравним время выполнения разных методов при решении одной и той же задачи — создания столбца 'D' на основе нескольких условий:

# Метод 1: np.where() вложенный
def method_np_where(df):
df = df.copy()
df['D'] = np.where(
(df['A'] > 80) & (df['B'] == 'X'),
'Group1',
np.where(
((df['A'] > 50) & (df['A'] <= 80)) | ((df['B'] == 'Y') & (df['C'] > 0)),
'Group2',
'Group3'
)
)
return df

# Метод 2: np.select()
def method_np_select(df):
df = df.copy()
conditions = [
(df['A'] > 80) & (df['B'] == 'X'),
((df['A'] > 50) & (df['A'] <= 80)) | ((df['B'] == 'Y') & (df['C'] > 0))
]
choices = ['Group1', 'Group2']
df['D'] = np.select(conditions, choices, default='Group3')
return df

# Метод 3: loc[]
def method_loc(df):
df = df.copy()
df['D'] = 'Group3' # значение по умолчанию
df.loc[(df['A'] > 80) & (df['B'] == 'X'), 'D'] = 'Group1'
df.loc[
((df['A'] > 50) & (df['A'] <= 80)) | ((df['B'] == 'Y') & (df['C'] > 0)), 
'D'
] = 'Group2'
return df

# Метод 4: apply() с lambda
def method_apply_lambda(df):
df = df.copy()
df['D'] = df.apply(
lambda row: 'Group1' if row['A'] > 80 and row['B'] == 'X'
else 'Group2' if ((row['A'] > 50 and row['A'] <= 80) or 
(row['B'] == 'Y' and row['C'] > 0))
else 'Group3',
axis=1
)
return df

# Метод 5: apply() с именованной функцией
def categorize_row(row):
if row['A'] > 80 and row['B'] == 'X':
return 'Group1'
elif ((row['A'] > 50 and row['A'] <= 80) or 
(row['B'] == 'Y' and row['C'] > 0)):
return 'Group2'
else:
return 'Group3'

def method_apply_func(df):
df = df.copy()
df['D'] = df.apply(categorize_row, axis=1)
return df

# Измеряем время выполнения каждого метода
results = {}
for method_name, method_func in [
('np.where()', method_np_where),
('np.select()', method_np_select),
('loc[]', method_loc),
('apply()+lambda', method_apply_lambda),
('apply()+func', method_apply_func)
]:
_, execution_time = time_execution(method_func, test_df)
results[method_name] = execution_time

Результаты тестов показывают, что векторизованные операции (np.where() и np.select()) значительно превосходят по скорости методы, работающие со строками по отдельности (apply()). Это особенно заметно на больших датасетах, где разница может составлять 20-50 раз!

Ключевые выводы из сравнения:

• np.where() и np.select() — лучший выбор для обработки больших объемов данных, когда важна производительность.
• loc[] — хороший компромисс между производительностью и читаемостью для средних объемов данных.
• apply() с lambda или функциями — предпочтительно для сложной логики, которую трудно выразить векторизованно, или когда производительность не критична.

Интересно, что метод .loc[] показывает промежуточные результаты. Это связано с тем, что pandas оптимизирует некоторые операции фильтрации, но всё равно выполняет больше работы, чем чистые numpy-методы.

Для реальных задач рекомендую следовать этой стратегии:

1. Сначала попробуйте использовать np.select() или np.where() для максимальной производительности.
2. Если логика слишком сложная для векторизации или требует особой обработки, используйте .loc[].
3. Только когда предыдущие подходы неприменимы из-за очень специфичных требований, переходите к .apply().

Помните, что производительность — это только один аспект. Читаемость, поддерживаемость и корректность кода не менее важны, особенно в командной работе или долгосрочных проектах.

Выбор метода добавления столбцов в pandas с множественными условиями — это баланс между производительностью, читаемостью и гибкостью. Для больших датасетов стоит отдавать предпочтение векторизованным операциям np.where() и np.select(), которые могут быть в 20-50 раз быстрее, чем функциональные подходы. Если же требуется выразить сложную бизнес-логику с множеством условий и вычислений, использование .loc[] или .apply() предоставит необходимую гибкость. Не забывайте, что хороший код — это не только быстрый код, но и понятный для ваших коллег и будущего вас. Экспериментируйте, тестируйте производительность на ваших реальных данных и выбирайте подход, который лучше всего соответствует конкретной ситуации.

Читайте также

• Pandas для Python: мощный инструмент анализа и обработки данных
• Python и Google Sheets API: автоматизация работы с таблицами
• Интеграция PyCharm и ClickUp: автоматизация задач разработки
• Как добавить строку в DataFrame pandas: 5 эффективных методов
• Как установить Anaconda и Jupyter Notebook для работы с данными
• PyTorch и TensorFlow: выбор фреймворка для задач машинного обучения
• Автоматизация Google Таблиц через Python: пишем код, экономим время
• Как использовать значения словарей в Python: полное руководство
• Асинхронное программирование в Python: скорость, возможности, практика
• Онлайн интерпретаторы Python: 8 лучших решений для кода везде