Kumulierte Summe mit group_by ⇐ Python

[Anonymous]

Angenommen, ich habe den folgenden DataFrame der Anzahl der Geburten in jedem Bundesstaat in jedem Jahr:

Code: Select all

df = pl.DataFrame(
{
"state": ["CA", "CA", "CA", "TX", "TX", "TX"],
"year": [2010, 2011, 2012, 2010, 2011, 2012],
"births": [10, 20, 30, 40, 50, 60],
}
)

Angenommen, die Bevölkerung im Jahr 2009 beträgt in allen Bundesstaaten 0.
Angenommen, die Zahl der Todesfälle beträgt in allen Bundesstaaten auch für alle Jahre 0.
Ich möchte die Bevölkerung in jedem Bundesstaat für jedes Jahr berechnen, indem ich die kumulative Geburtensumme der einzelnen Bundesstaaten für jedes Jahr verwende.
Ich habe mir Folgendes ausgedacht:

Code: Select all

df.group_by("state").agg(pl.col("births").cum_sum().alias("pop")).explode(
"pop"
).with_columns(year=df["year"])

Dazu ist jedoch das erneute Hinzufügen der Jahresspalte aus dem vorherigen Datensatz erforderlich, was nicht sehr sauber ist.
Gibt es einen besseren, saubereren Weg?
Die Dokumentation für cum_sum enthält kein Beispiel mit group_by