Векторизовать условное присвоение в pandas dataframe

Если у меня есть фрейм данных df со столбцом x и я хочу создать столбец y на основе значений x, используя это в псевдокоде:

 if df['x'] <-2 then df['y'] = 1 
 else if df['x'] > 2 then df['y']= -1 
 else df['y'] = 0

Как бы я этого достиг? Я предполагаю, что np.where - лучший способ сделать это, но не уверен, как правильно его кодировать.

ОТВЕТЫ

Ответ 1

Одним простым способом было бы сначала присвоить значение по умолчанию, а затем выполнить 2 loc вызова:

In [66]:

df = pd.DataFrame({'x':[0,-3,5,-1,1]})
df
Out[66]:
   x
0  0
1 -3
2  5
3 -1
4  1

In [69]:

df['y'] = 0
df.loc[df['x'] < -2, 'y'] = 1
df.loc[df['x'] > 2, 'y'] = -1
df
Out[69]:
   x  y
0  0  0
1 -3  1
2  5 -1
3 -1  0
4  1  0

Если вы хотите использовать np.where, вы можете сделать это с помощью вложенного np.where:

In [77]:

df['y'] = np.where(df['x'] < -2 , 1, np.where(df['x'] > 2, -1, 0))
df
Out[77]:
   x  y
0  0  0
1 -3  1
2  5 -1
3 -1  0
4  1  0

Итак, здесь мы определяем первое условие, где x меньше, чем -2, return 1, тогда мы имеем еще один np.where, который проверяет другое условие, где x больше 2 и возвращает -1, в противном случае возвращает 0

<сильные > тайминги

In [79]:

%timeit df['y'] = np.where(df['x'] < -2 , 1, np.where(df['x'] > 2, -1, 0))

1000 loops, best of 3: 1.79 ms per loop

In [81]:

%%timeit
df['y'] = 0
df.loc[df['x'] < -2, 'y'] = 1
df.loc[df['x'] > 2, 'y'] = -1

100 loops, best of 3: 3.27 ms per loop

Итак, для этого набора данных образца метод np.where в два раза быстрее

Ответ 2

Это хороший вариант использования для pd.cut, где вы определяете диапазоны и на основе этих ranges вы можете назначить labels:

df['y'] = pd.cut(df['x'], [-np.inf, -2, 2, np.inf], labels=[1, 0, -1], right=False)

Выход

   x  y
0  0  0
1 -3  1
2  5 -1
3 -1  0
4  1  0