Drill — Solutions Hidden by Default

Pandas Problems

Ten drillable LeetCode-style problems solved in pandas, with multiple approaches per problem and discussion of when to push to SQL/DuckDB at scale.

🎯 10 problems ⏱ ~15 min each 📊 Progress saved
0 / 10 practiced

Problems

P1. Customers with monotonically increasing weekly revenue

Given sessions(customer_id, started_at, revenue), return the customer_ids whose weekly revenue strictly increased every week for the last 4 weeks.

Show solution
monotone_increasing.py
import pandas as pd

def monotone_customers(sessions: pd.DataFrame) -> list[int]:
    sessions = sessions.copy()
    sessions['week'] = pd.to_datetime(sessions['started_at']).dt.to_period('W').dt.start_time
    weekly = (sessions.groupby(['customer_id', 'week'], as_index=False)['revenue'].sum()
                       .sort_values(['customer_id', 'week']))
    # Last 4 weeks per customer
    weekly['rk'] = weekly.groupby('customer_id')['week'].rank(method='first', ascending=False)
    last4 = weekly[weekly['rk'] <= 4].sort_values(['customer_id', 'week'])
    # Check strict increase
    last4['delta'] = last4.groupby('customer_id')['revenue'].diff()
    grouped = last4.groupby('customer_id').agg(
        n_weeks=('week', 'count'),
        all_increasing=('delta', lambda s: (s.dropna() > 0).all() and s.notna().sum() == 3),
    )
    return grouped[(grouped['n_weeks'] == 4) & (grouped['all_increasing'])].index.tolist()

Discussion: the diff() + check pattern is the pandas equivalent of LAG + comparison in SQL. At scale, this is much faster as SQL on a warehouse — name that as the alternative.

P2. Top-K per group

Given a DataFrame of (customer_id, gpu_id, revenue), return the top 3 GPUs by revenue per customer.

Show solution
top_k_per_group.py
import pandas as pd

# Approach 1: sort + groupby + head
def top_k_v1(df, k=3):
    return (df.sort_values(['customer_id', 'revenue'], ascending=[True, False])
              .groupby('customer_id')
              .head(k))

# Approach 2: rank + filter
def top_k_v2(df, k=3):
    df = df.copy()
    df['rk'] = df.groupby('customer_id')['revenue'].rank(method='first', ascending=False)
    return df[df['rk'] <= k].drop(columns='rk')

# Approach 3: nlargest per group — pandas idiomatic
def top_k_v3(df, k=3):
    return df.groupby('customer_id', group_keys=False).apply(
        lambda g: g.nlargest(k, 'revenue'), include_groups=False
    )

Tradeoffs: v1 is fastest and clearest. v2 lets you keep all tied rows if you switch to method='min'. v3 reads cleanest but is slowest (Python lambda). The senior move: mention all three.

P3. Sessionize event stream

Given events(customer_id, event_at), group events into sessions where a gap > 30 minutes starts a new session. Return (customer_id, session_id, n_events).

Show solution
sessionize.py
import pandas as pd

def sessionize(events: pd.DataFrame, gap_minutes: int = 30) -> pd.DataFrame:
    events = events.sort_values(['customer_id', 'event_at']).copy()
    events['event_at'] = pd.to_datetime(events['event_at'])
    prev = events.groupby('customer_id')['event_at'].shift(1)
    is_new = (prev.isna()) | ((events['event_at'] - prev).dt.total_seconds() / 60 > gap_minutes)
    events['session_id'] = is_new.astype(int).groupby(events['customer_id']).cumsum()
    return events.groupby(['customer_id', 'session_id']).size().reset_index(name='n_events')

Same trick as the SQL version — flag new-session boundaries, then cumsum within each customer to assign session IDs.

P4. 7-day rolling DAU

Given events(user_id, event_date), compute the 7-day rolling DAU per day.

Show solution
rolling_dau.py
import pandas as pd

def rolling_7d_dau(events: pd.DataFrame) -> pd.DataFrame:
    events['event_date'] = pd.to_datetime(events['event_date'])
    # Per (user, day) we need to know "did user appear in [day-6, day]"
    # Naive nested loop is O(n*days). Use a clever set-membership approach via groupby+rolling
    # Daily distinct users
    daily_users = events.groupby('event_date')['user_id'].apply(set).rename('users')
    days = pd.date_range(daily_users.index.min(), daily_users.index.max(), freq='D')
    daily_users = daily_users.reindex(days, fill_value=set())
    # Rolling union over 7-day window
    result = []
    from collections import deque
    window = deque(maxlen=7)
    for d, users in daily_users.items():
        window.append(users)
        union = set().union(*window)
        result.append((d, len(union)))
    return pd.DataFrame(result, columns=['day', 'dau_7d'])

The discussion to volunteer: "pandas .rolling() doesn't support distinct-count out of the box. At scale, this is a SQL job — COUNT(DISTINCT user_id) OVER (ORDER BY day ROWS 6 PRECEDING) works in some dialects, or HyperLogLog approximations in BigQuery/Redshift. I wrote the deque version because it's O(n × avg_dau) instead of O(n × days × avg_dau)."

P5. Detect anomalous daily revenue

Given (day, revenue), flag days where revenue is > 3 standard deviations from the trailing 28-day mean (excluding the day itself).

Show solution
anomaly_3sigma.py
import pandas as pd

def flag_anomalies(df: pd.DataFrame, window: int = 28, k: float = 3) -> pd.DataFrame:
    df = df.sort_values('day').reset_index(drop=True).copy()
    # Trailing window, excluding the current row → shift first, then rolling
    history = df['revenue'].shift(1)
    df['mu'] = history.rolling(window=window, min_periods=7).mean()
    df['sigma'] = history.rolling(window=window, min_periods=7).std()
    df['anomalous'] = (df['revenue'] - df['mu']).abs() > k * df['sigma']
    return df

The discussion to volunteer: 3σ is a thin definition. Real anomaly detection would use robust stats (MAD), seasonal decomposition, or per-weekday baselines because GPU revenue has weekday cycles. State this unprompted.

P6. Two-sum on transaction amounts

Given a list of GPU session revenues and a target T, return any pair of session indices whose revenues sum to T.

Show solution
two_sum.py
def two_sum(revenues: list[float], target: float) -> tuple[int, int] | None:
    seen: dict[float, int] = {}
    for i, r in enumerate(revenues):
        complement = target - r
        if complement in seen:
            return (seen[complement], i)
        seen[r] = i
    return None

O(n) with O(n) space. Classic. Worth knowing the follow-up: "all pairs that sum to T" → store all indices per value in a dict-of-lists.

P7. Customer cohort retention as DataFrame

Given signups (customer_id, signed_up_at) and sessions (customer_id, event_at), return a wide DataFrame: rows are signup-week cohorts, columns are weeks_since (0–8), values are retention rate.

Show solution
cohort_retention.py
import pandas as pd

def cohort_retention(signups, sessions, max_weeks=8):
    signups = signups.copy()
    signups['cohort_week'] = pd.to_datetime(signups['signed_up_at']).dt.to_period('W').dt.start_time
    cohort_size = signups.groupby('cohort_week').size().rename('size')

    sessions = sessions.merge(signups[['customer_id', 'cohort_week']], on='customer_id', how='inner')
    sessions['active_week'] = pd.to_datetime(sessions['event_at']).dt.to_period('W').dt.start_time
    sessions['weeks_since'] = ((sessions['active_week'] - sessions['cohort_week']).dt.days // 7)

    retained = (sessions[(sessions['weeks_since'] >= 0) & (sessions['weeks_since'] <= max_weeks)]
                .groupby(['cohort_week', 'weeks_since'])['customer_id']
                .nunique()
                .rename('retained')
                .reset_index())
    retained = retained.merge(cohort_size, on='cohort_week')
    retained['rate'] = retained['retained'] / retained['size']
    return retained.pivot(index='cohort_week', columns='weeks_since', values='rate')

P8. Carry forward last non-null plan

Given (customer_id, event_at, plan) where plan can be null on some rows, fill each null with the customer's most recent prior non-null plan.

Show solution
ffill_per_group.py
import pandas as pd

def forward_fill_plan(df: pd.DataFrame) -> pd.DataFrame:
    df = df.sort_values(['customer_id', 'event_at']).copy()
    df['plan'] = df.groupby('customer_id')['plan'].ffill()
    return df

One-liner: .ffill() after a groupby does the right thing. The trap is forgetting to sort first.

P9. Pivot: hours per tier per customer

Given (customer_id, tier, gpu_hours), return a wide table with one row per customer and one column per tier (a100, h100, rtx-4090). Missing combinations should be 0.

Show solution
pivot_hours.py
import pandas as pd

def hours_by_tier(df: pd.DataFrame) -> pd.DataFrame:
    return df.pivot_table(
        index='customer_id', columns='tier', values='gpu_hours',
        aggfunc='sum', fill_value=0
    ).reset_index()

One call. Mention the SQL equivalent: SUM(CASE WHEN tier = 'a100' THEN gpu_hours END). Show you can do both.

P10. Find the first session after a price change

Given price_changes(tier, changed_at, new_price) and sessions(customer_id, tier, started_at), for each session find the most recent price change for that tier as of the session start.

Show solution
asof_price.py
import pandas as pd

def attach_price(sessions, price_changes):
    sessions = sessions.sort_values('started_at').copy()
    price_changes = price_changes.sort_values('changed_at').copy()
    return pd.merge_asof(
        sessions, price_changes,
        left_on='started_at', right_on='changed_at',
        by='tier',
        direction='backward',
    )

The discussion to volunteer: merge_asof with by='tier' is the point-in-time-join idiom. Without it, you'd be writing nested loops or a horrible cross-join. Senior tell: knowing this exists and reaching for it.

Drill protocol

How to use this page

Enable drill mode. Read each problem. 12-minute timer. Code in a notebook with example data. Talk through your approach + complexity out loud before writing. When solving, write the vectorized version first; only reach for apply with a comment. Reveal, compare, mark practiced. Aim for 7/10 cleared before the screen.