Contents

Python Pandas For Your Grandpa - 3.11 Challenge: Party Time

2021-02-01 389 words 2 minutes

Contents

Setup

Whenever your friends John and Judy visit you together, y’all have a party. Given a DataFrame with 10 rows representing the next 10 days of your schedule and whether John and Judy are scheduled to make an appearance, insert a new column called days_til_party that indicates how many days until the next party relative to each day.

days_til_party should be 0 on days where a party is occurs, 1 on days where a party doesn’t occur but will occur the next day, etc.

import numpy as np
import pandas as pd

generator = np.random.default_rng(123)
df = pd.DataFrame({
    'john': generator.choice([True, False], size=10, replace=True),
    'judy': generator.choice([True, False], size=10, replace=True)
})
print(df)
##     john   judy
## 0   True   True
## 1  False  False
## 2  False   True
## 3   True  False
## 4  False   True
## 5   True   True
## 6   True  False
## 7   True  False
## 8   True  False
## 9   True  False

Solution

party = df.john & df.judy
grps = party.iloc[::-1].cumsum()
df['days_til_party'] = party.groupby(grps).cumcount(ascending=False)
df.loc[(party.loc[party].index[-1] + 1):, 'days_til_party'] = pd.NA
print(df)
##     john   judy days_til_party
## 0   True   True              0
## 1  False  False              4
## 2  False   True              3
## 3   True  False              2
## 4  False   True              1
## 5   True   True              0
## 6   True  False           <NA>
## 7   True  False           <NA>
## 8   True  False           <NA>
## 9   True  False           <NA>

Course Curriculum