web-dev-qa-db-fra.com

vectoriser l'affectation conditionnelle dans pandas dataframe

Si j'ai un dataframe df avec la colonne x et que je veux créer la colonne y sur la base des valeurs de x en utilisant ceci dans le pseudo code:

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

Comment pourrais-je y parvenir? Je suppose np.where est le meilleur moyen de le faire, mais vous ne savez pas comment le coder correctement.

26
azuric

Une méthode simple serait d'attribuer d'abord la valeur par défaut, puis d'effectuer 2 appels 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

Si vous souhaitez utiliser np.where, Vous pouvez le faire avec un np.where Imbriqué:

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

Donc, ici, nous définissons la première condition comme où x est inférieur à -2, retournons 1, puis nous avons un autre np.where Qui teste l'autre condition où x est supérieur à 2 et renvoie -1, sinon retourne 0

horaires

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

Ainsi, pour cet exemple de jeu de données, la méthode np.where Est deux fois plus rapide

26
EdChum

Ceci est un bon cas d'utilisation pour pd.cut où vous définissez des plages et sur la base de ces ranges vous pouvez affecter labels:

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

Sortie

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