Добавление столбца в pandas по условиям

Введение

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

Основы pandas и DataFrame

pandas — это библиотека Python для анализа данных, которая предоставляет мощные инструменты для работы с табличными данными. Основной структурой данных в pandas является DataFrame, который представляет собой таблицу с метками строк и столбцов. DataFrame позволяет легко манипулировать данными, фильтровать их, добавлять новые столбцы и выполнять множество других операций.

import pandas as pd

# Создание простого DataFrame
data = {
    'Name': ['Alice', 'Bob', 'Charlie'],
    'Age': [25, 30, 35],
    'Salary': [50000, 60000, 70000]
}

df = pd.DataFrame(data)
print(df)

DataFrame в pandas можно создать из различных источников данных, включая CSV-файлы, базы данных и даже другие DataFrame. В приведенном выше примере мы создали DataFrame из словаря Python, где ключи представляют собой имена столбцов, а значения — списки данных.

Добавление столбца по одному условию

Для добавления нового столбца на основе одного условия можно использовать метод loc. Например, добавим столбец Senior, который будет содержать значение True, если возраст больше 30 лет, и False в противном случае.

df['Senior'] = df['Age'] > 30
print(df)

Результат будет следующим:

      Name  Age  Salary  Senior
0    Alice   25   50000   False
1      Bob   30   60000   False
2  Charlie   35   70000    True

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

Добавление столбца по нескольким условиям

Для более сложных условий можно использовать логические операторы и функции numpy. Рассмотрим пример, где мы добавим столбец Category, который будет содержать значение 'High', если зарплата больше 60000 и возраст больше 30, и 'Low' в противном случае.

import numpy as np

conditions = [
    (df['Salary'] > 60000) & (df['Age'] > 30),
    (df['Salary'] <= 60000) | (df['Age'] <= 30)
]

choices = ['High', 'Low']

df['Category'] = np.select(conditions, choices)
print(df)

Результат будет следующим:

      Name  Age  Salary  Senior Category
0    Alice   25   50000   False      Low
1      Bob   30   60000   False      Low
2  Charlie   35   70000    True     High

Функция np.select из библиотеки numpy позволяет легко управлять множественными условиями и значениями. Она принимает два аргумента: список условий и список значений, которые будут присвоены, если условие выполнено.

Примеры и практические советы

Пример 1: Добавление столбца с категорией по возрасту

Добавим столбец AgeGroup, который будет содержать значение 'Young', если возраст меньше 30, 'Middle', если возраст от 30 до 40, и 'Old', если возраст больше 40.

conditions = [
    (df['Age'] < 30),
    (df['Age'] >= 30) & (df['Age'] <= 40),
    (df['Age'] > 40)
]

choices = ['Young', 'Middle', 'Old']

df['AgeGroup'] = np.select(conditions, choices)
print(df)

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

Пример 2: Добавление столбца с уровнем зарплаты

Добавим столбец SalaryLevel, который будет содержать значение 'Low', если зарплата меньше 55000, 'Medium', если зарплата от 55000 до 65000, и 'High', если зарплата больше 65000.

conditions = [
    (df['Salary'] < 55000),
    (df['Salary'] >= 55000) & (df['Salary'] <= 65000),
    (df['Salary'] > 65000)
]

choices = ['Low', 'Medium', 'High']

df['SalaryLevel'] = np.select(conditions, choices)
print(df)

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

Пример 3: Добавление столбца с комбинированными условиями

Добавим столбец Status, который будет содержать значение 'Senior High', если возраст больше 30 и зарплата больше 60000, 'Senior Low', если возраст больше 30 и зарплата меньше или равна 60000, и 'Junior' в противном случае.

conditions = [
    (df['Age'] > 30) & (df['Salary'] > 60000),
    (df['Age'] > 30) & (df['Salary'] <= 60000),
    (df['Age'] <= 30)
]

choices = ['Senior High', 'Senior Low', 'Junior']

df['Status'] = np.select(conditions, choices)
print(df)

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

Практические советы

1. Используйте numpy для сложных условий: Функция np.select позволяет легко управлять множественными условиями и значениями.
2. Оптимизируйте производительность: Если у вас большой DataFrame, старайтесь минимизировать количество операций, выполняемых над данными. Например, объединяйте условия и используйте векторизованные операции.
3. Проверяйте данные: Перед добавлением новых столбцов убедитесь, что ваши условия корректны и данные не содержат ошибок. Это поможет избежать неожиданных результатов и ошибок в анализе.
4. Используйте методы pandas для фильтрации данных: Методы query и filter могут быть полезны для фильтрации данных перед добавлением новых столбцов.
5. Документируйте свои шаги: Добавляйте комментарии к коду, чтобы другие пользователи могли легко понять вашу логику и условия.

Заключение

Добавление столбцов в pandas по условиям — это мощный инструмент для анализа данных. Используя методы, описанные в этой статье, вы сможете легко создавать новые столбцы на основе различных условий. Надеюсь, что эти примеры и советы помогут вам в вашей работе с pandas. 😉