Панды добавить столбец со значением на основе условия на основе других столбцов

Question

У меня есть следующий кадр данных панд:

import pandas as pd
import numpy as np

d = {'age' : [21, 45, 45, 5],
     'salary' : [20, 40, 10, 100]}

df = pd.DataFrame(d)

и хотел бы добавить дополнительный столбец под названием «is_rich», который фиксирует, богат ли человек в зависимости от его / ее зарплаты. Я нашел несколько способов сделать это:

# method 1
df['is_rich_method1'] = np.where(df['salary']>=50, 'yes', 'no')

# method 2
df['is_rich_method2'] = ['yes' if x >= 50 else 'no' for x in df['salary']]

# method 3
df['is_rich_method3'] = 'no'
df.loc[df['salary'] > 50,'is_rich_method3'] = 'yes'

в результате:

Однако я не понимаю, какой предпочтительный способ. Все ли методы одинаково хороши в зависимости от вашего приложения?

Answer 1

Используйте timeits, Люк!

Заключение
Постижения списков лучше всего работают с меньшими объемами данных, поскольку они требуют очень мало накладных расходов, даже если они не векторизованы. OTOH, на больших данных, loc и numpy.where работают лучше - векторизация побеждает день.

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

Тем не менее, одно из достоинств этого факта заключается в том, что списочные представления довольно конкурентоспособны - они реализованы в C и высоко оптимизированы для производительности.

---

Код для сравнения, для справки . Вот функции по времени:

def numpy_where(df):
  return df.assign(is_rich=np.where(df['salary'] >= 50, 'yes', 'no'))

def list_comp(df):
  return df.assign(is_rich=['yes' if x >= 50 else 'no' for x in df['salary']])

def loc(df):
  df = df.assign(is_rich='no')
  df.loc[df['salary'] > 50, 'is_rich'] = 'yes'
  return df