[FIXED] pandas groupby().head(n) wobei n eine Funktion der Gruppenbezeichnung ist

Ausgabe

Ich habe einen Datenrahmen und möchte nach einer Spalte gruppieren und den Kopf jeder Gruppe nehmen, aber ich möchte, dass die Tiefe des Kopfs durch eine Funktion der Gruppenbezeichnung definiert wird. Wenn da nicht die variablen Gruppengrößen wären, könnte ich locker df.groupby('label').head(n). Ich kann mir eine Lösung vorstellen, die das Durchlaufen von df['label'].unique(), das Aufteilen des Datenrahmens und das Erstellen eines neuen beinhaltet, aber ich befinde mich in einem Kontext, in dem ich ziemlich leistungsempfindlich bin, daher möchte ich diese Art von Iteration nach Möglichkeit vermeiden.

Hier ist ein Beispiel für einen Datenrahmen:

  label   values
0  apple       7
1  apple       5
2  apple       4
3    car       9
4    car       6
5    dog       5
6    dog       3
7    dog       2
8    dog       1

und Code für mein Beispiel-Setup:

import pandas as pd
df = pd.DataFrame({'label': ['apple', 'apple', 'apple', 'car', 'car', 'dog', 'dog', 'dog', 'dog'],
          'values': [7, 5, 4, 9, 6, 5, 3, 2 ,1]})
def depth(label):
    if label == 'apple': return 1
    elif label == 'car': return 2
    elif label == 'dog': return 3

Meine gewünschte Ausgabe ist ein Datenrahmen mit der Anzahl der Zeilen aus jeder Gruppe, die durch diese Funktion definiert wird:

   label  values
0  apple       7
3    car       9
4    car       6
5    dog       5
6    dog       3
7    dog       2

Lösung

Ich würde hier ein Wörterbuch verwenden und <group>.namein verwenden groupby.apply:

depth = {'apple': 1, 'car': 2, 'dog': 3}

out = (df.groupby('label', group_keys=False)
         .apply(lambda g: g.head(depth.get(g.name, 0)))
       )

NB. Wenn Sie eine Funktion wirklich brauchen, können Sie dasselbe mit einem Funktionsaufruf tun. Stellen Sie sicher, dass Sie in jedem Fall einen Wert zurückgeben.

Alternative Option mit groupby.cumcountund boolescher Indizierung:

out = df[df['label'].map(depth).gt(df.groupby('label').cumcount())]

Ausgang:

   label  values
0  apple       7
3    car       9
4    car       6
5    dog       5
6    dog       3
7    dog       2


Beantwortet von –
mozway


Antwort geprüft von –
David Marino (FixError Volunteer)

0 Shares:
Leave a Reply

Your email address will not be published. Required fields are marked *

You May Also Like