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):
    top = df.groupby('customer_id')['revenue'].nlargest(k)
    return df.loc[top.index.get_level_values(-1)]

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 (...) is rejected by PostgreSQL and many other engines, so join each day to the six days before it and count distinct users, or use 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.

✓ Check yourself

Answer from memory before opening each one.

  1. In the sessionization problem, events for one customer are at 10:00, 10:20 and 11:00 with a 30-minute gap. Trace the shifted column, the new-session flags and the session ids.

    Show answer

    The shift gives missing, 10:00, 10:20, so the flags are true, false, true because a missing previous value starts a session and the 40-minute gap exceeds 30. The cumulative sum gives ids 1, 1, 2. The data must be sorted by customer and time first. See the problems.

  2. The anomaly detector shifts revenue by one day before the rolling mean and standard deviation. What breaks if you leave the shift out?

    Show answer

    The current day sits inside its own baseline, so a spike raises the mean and the deviation and partly hides itself. With the shift, the baseline is the prior days only, as the problem requires. See the problems.

  3. A forward fill of plan is run per customer on rows that are not sorted by time. What does it produce, and what if you forget to group?

    Show answer

    It fills each missing value from whichever row happens to precede it in the frame, which may be a later event. Without the group, a customer's missing plan takes the previous customer's last plan. Sort by customer and time, then group, then fill. See the problems.

  4. Why can't a rolling window compute seven-day distinct users directly, and what does the deque solution trade?

    Show answer

    Rolling supports aggregates like mean and sum that update incrementally, not a distinct count over a window of values. The deque holds sets of users per day and unions them, which costs memory and set operations per day but avoids rescanning. At scale you would do it in the warehouse. See the problems.

  5. For top-3 GPUs per customer, a customer has only two GPUs, and another has two tied for third place. What do the head and rank approaches return in each case?

    Show answer

    The first customer returns two rows from either approach, since there are no more. For the tie, head with a sort returns exactly three rows and drops one of the tied pair arbitrarily. Rank with method min keeps both. See the problems.

Takeaway

  • Sessionization in pandas mirrors SQL: shift to compare with the previous event, flag boundaries, then cumulative sum within each customer.
  • Trailing baselines must exclude the current row, so shift before you roll.
  • Forward fill is only correct after sorting by time and grouping by entity.
  • Rolling cannot do distinct counts, so choose a set-based loop for small data and the warehouse for large.
  • Top-k per group depends on tie handling, so state which behaviour you chose.
  • merge_asof with by is the point-in-time join idiom.

The last piece is the domain the company works in, and the next chapter gives you its vocabulary. → GPU Compute Domain

Further reading