Manière efficace d'appliquer un masque booléen multiple pour définir des valeurs dans une colonne en utilisant des pandas

Question

J'ai plusieurs expressions qui sélectionnent certaines lignes dans un bloc de données (df) et retournent plusieurs tableaux booléens, masques si vous le souhaitez. Je souhaite ajouter une nouvelle colonne à ce bloc de données et le remplir avec des entiers triviaux (indicateurs d'appartenance) basés sur ces masques. Dans l'approche la plus simple j'utiliser ce code:

df.assign(membership=3) 
df[mask1].membership = 0 
df[mask2].membership = 1 
df[mask3].membership = 2 

Cela devrait donner une colonne (df["membership"]) comme suit:

1 0 
2 0 
3 1 
4 1 
5 3 
Name: membership, dtype: int64 

Les masques sont créés par la fonction (en utilisant des colonnes supplémentaires de la trame de données pour créer la sortie) comme un dictionnaire, qui ressemble à ceci:

maskDict = {0:mask1, 1:mask2, 2:mask3} 

où mask(1..3) sont des tableaux de taille égale contenantJe suis conscient qu'il ne sera peut-être pas nécessaire d'améliorer le code. Cependant, en comparant plus de conditions je finirai avec plus de tableaux booléens dans le dictionnaire.

Une solution sur place qui applique la fonction à la trame de données et définit ces valeurs serait également appréciée.

Answer 1

Utilisation numpy.select:

df["membership"] = np.select([mask1, mask2, mask3], [0,1,2], default=3) 

Si l'entrée est dict:

df["membership1"] = np.select(list(maskDict.values()), list(maskDict.keys()), default=3)