Section C · Python & pandas

Pandas Fundamentals

The pandas idioms LeetCode-style screens reach for — groupby, merge, time-series, vectorization, and the SQL↔pandas translation that wins the "could you also do this in SQL?" follow-up.

The interview approach

  1. Restate the problem. Ask about input shape (DataFrame vs list of dicts), expected output shape, edge cases (empty input, NaNs, ties).
  2. State your approach: "I'm going to groupby A, aggregate B, then merge back." Get a nod.
  3. Write the vectorized version first. Reach for .apply() only with a comment about why it's the right call.
  4. Mention the alternative: "Could also push to SQL or use DuckDB on this DataFrame at scale." This is the senior tell.
  5. Trace 2 example rows. Catches off-by-ones.

Groupby patterns

Basic aggregates

multi-aggregate groupby
import pandas as pd

# Multiple aggregations per group, named outputs
result = df.groupby('customer_id').agg(
    n_sessions=('session_id', 'count'),
    total_revenue=('revenue', 'sum'),
    first_session=('started_at', 'min'),
    last_session=('started_at', 'max'),
)

# Aggregate different columns differently in one call
result = df.groupby('tier').agg({
    'revenue': 'sum',
    'session_id': 'nunique',
    'started_at': ['min', 'max'],
})

Filter groups (HAVING-equivalent)

groups with >= N rows
# Keep customers with at least 5 sessions
result = df.groupby('customer_id').filter(lambda g: len(g) >= 5)

# Equivalent, faster — avoids the Python lambda
counts = df.groupby('customer_id').size()
keep = counts[counts >= 5].index
result = df[df['customer_id'].isin(keep)]

Transform: add a group-level value back to every row

group share of total
# Each row's revenue as a fraction of the customer's total
df['share_of_total'] = df['revenue'] / df.groupby('customer_id')['revenue'].transform('sum')

Pivot

long-to-wide
wide = df.pivot_table(
    index='customer_id',
    columns='tier',
    values='gpu_hours',
    aggfunc='sum',
    fill_value=0,
)

Merge gotchas

  • Default how='inner' — always specify how explicitly. 'left' is what you usually want for "enrich the left table."
  • Row duplication: if your join key isn't unique on the right side, you'll get more rows than you started with. Always check row count before and after.
  • Column-name collisions: when both sides have a column not in the join key, pandas appends _x and _y. Pre-rename or use suffixes.
  • validate argument: df.merge(other, on='key', how='left', validate='m:1') raises if the assumption is wrong. Use it.
safe merge
enriched = sessions.merge(
    customers[['customer_id', 'plan', 'country']],
    on='customer_id',
    how='left',
    validate='m:1',  # raises if customer_id isn't unique in customers
)
assert len(enriched) == len(sessions)  # row count sanity

The asof merge

For "what was the value of X at the time of event Y" — point-in-time joins.

price-at-session-time
sessions = sessions.sort_values('started_at')
price_history = price_history.sort_values('changed_at')

result = pd.merge_asof(
    sessions, price_history,
    left_on='started_at', right_on='changed_at',
    by='tier',                  # match within the same tier
    direction='backward',       # most-recent price at or before started_at
    tolerance=pd.Timedelta('30 days'),  # don't match older than 30 days
)

Window-function equivalents

SQLPandas
ROW_NUMBER() OVER (PARTITION BY a ORDER BY b)df.sort_values('b').groupby('a').cumcount() + 1
RANK() OVER (...)df.groupby('a')['b'].rank(method='min')
DENSE_RANK() OVER (...)df.groupby('a')['b'].rank(method='dense')
LAG(b) OVER (PARTITION BY a ORDER BY t)df.sort_values('t').groupby('a')['b'].shift(1)
LEAD(b) OVER (...)df.sort_values('t').groupby('a')['b'].shift(-1)
SUM(b) OVER (PARTITION BY a)df.groupby('a')['b'].transform('sum')
SUM(b) OVER (PARTITION BY a ORDER BY t) (running)df.sort_values('t').groupby('a')['b'].cumsum()
AVG(b) OVER (ORDER BY t ROWS 6 PRECEDING)df['b'].rolling(window=7, min_periods=1).mean()

Time-series ops

Resample

For "aggregate by week" / "fill missing days" / "downsample to hourly":

weekly resample
df['started_at'] = pd.to_datetime(df['started_at'])
weekly = (df.set_index('started_at')
            .groupby('customer_id')
            .resample('W-SUN')           # ISO weeks (Mon-Sun), labelled by the Sunday
            .agg({'revenue': 'sum', 'session_id': 'count'})
            .reset_index())

Fill missing dates

reindex on a date range
daily = df.set_index('day').reindex(
    pd.date_range(df['day'].min(), df['day'].max(), freq='D'),
    fill_value=0
)

Date math

  • (df['ended_at'] - df['started_at']).dt.total_seconds() / 3600 — gpu-hours.
  • df['day'] = df['started_at'].dt.normalize() — strip time component to midnight.
  • df['week'] = df['started_at'].dt.to_period('W').dt.start_time — week truncation.

Vectorization

Avoid .apply when possible

Pandas' .apply with a Python function is 10–100× slower than vectorized ops on big DataFrames. Reach for vectorization first.

slow vs vectorized
import numpy as np
# ✗ Slow
df['tier_label'] = df.apply(
    lambda r: 'high' if r['gpu_hours'] > 100 else 'low', axis=1)

# ✓ Vectorized
df['tier_label'] = np.where(df['gpu_hours'] > 100, 'high', 'low')

# Multi-branch: np.select
df['tier_label'] = np.select(
    [df['gpu_hours'] > 1000, df['gpu_hours'] > 100, df['gpu_hours'] > 0],
    ['enterprise', 'pro', 'free'],
    default='inactive'
)

When .apply is the right call

  • The function is genuinely complex (multi-line logic, branching, lookups against external data).
  • The DataFrame is small (<100k rows) and clarity matters more than speed.
  • You're applying along axis=0 (column-wise) to a small number of columns.

Mention DuckDB on big data

For "at scale" follow-ups: DuckDB can query a pandas DataFrame directly with SQL, often faster than pure pandas for groupby/join-heavy work:

duckdb on a DataFrame
import duckdb
result = duckdb.sql("""
    SELECT customer_id, COUNT(*) AS sessions, SUM(revenue) AS total
    FROM df
    WHERE started_at >= '2026-04-01'
    GROUP BY customer_id
""").to_df()

Mentioning this unprompted reads senior — it shows you know when pandas isn't the right tool.

SQL ↔ pandas translation

SQLPandas
SELECT col FROM tdf['col'] or df[['col']]
SELECT * FROM t WHERE x > 5df[df['x'] > 5]
SELECT * FROM t WHERE x IN (1,2,3)df[df['x'].isin([1,2,3])]
SELECT * FROM t WHERE x IS NULLdf[df['x'].isna()]
GROUP BY a HAVING COUNT(*) > 5df.groupby('a').filter(lambda g: len(g) > 5)
ORDER BY a DESCdf.sort_values('a', ascending=False)
LIMIT 10df.head(10)
DISTINCT a, bdf[['a','b']].drop_duplicates()
UNION ALLpd.concat([df1, df2], ignore_index=True)
JOIN ONdf1.merge(df2, on='id', how='inner')
COALESCE(a, b)df['a'].fillna(df['b'])
CASE WHEN ... THEN ... ENDnp.where / np.select

Edge cases & traps

The SettingWithCopyWarning

df[df['x'] > 0]['y'] = 99 may or may not modify the original DataFrame — pandas warns because the behavior depends on internal optimizations. The fix: df.loc[df['x'] > 0, 'y'] = 99.

NaN comparisons

NaN == NaN is False in pandas (and SQL). Use .isna() for tests; never == NaN.

Integer columns with NaN become floats

If you have a numeric column with any NaN, pandas casts the whole column to float. Use pd.Int64Dtype() (nullable integer) if you need to preserve integer semantics.

Object-dtype Series are slow

Strings stored as object are slow for many operations. pd.StringDtype() or pyarrow-backed strings are faster.

Sorting before groupby

Pandas' groupby sorts groups by default. If you've already sorted, pass sort=False to skip the redundant sort and save time.

✓ Check yourself

Answer from memory before opening each one.

  1. A left merge of 10,000 sessions with a customers table returns 14,200 rows. What happened, and which argument and check would have caught it?

    Show answer

    The customers key is not unique, so some sessions matched several rows and were repeated. Pass validate='m:1', which raises if the right side has duplicate keys, and assert the row count against the left table. See merge gotchas.

  2. A colleague uses df.apply(f, axis=1) on 5 million rows and it takes minutes. Why is it slow, and what do you change?

    Show answer

    apply with a Python function calls that function once per row, so the interpreter overhead dominates. Rewrite the logic with vectorized operations such as np.where or np.select. Keep apply only for logic too complex to vectorize on small data. See vectorization.

  3. Write the pandas equivalent of SUM(b) OVER (PARTITION BY a ORDER BY t). What goes wrong if you forget the sort?

    Show answer

    Sort by t, then group by a and take the cumulative sum of b. Without the sort, cumsum follows the existing row order, so the running total is in the wrong sequence and gives wrong values with no error. See window-function equivalents.

  4. A session at 10:00 for tier h100 meets price changes at 09:00 (4.00) and 10:30 (5.00). With merge_asof, direction backward and a 30-day tolerance, what price attaches? What if the last change was 45 days earlier?

    Show answer

    It attaches 4.00, the latest change at or before the session. If the last change was 45 days earlier, the tolerance rejects it and the price is NaN. Both frames must be sorted on their time keys, and by='tier' keeps matches within a tier. See the asof merge.

  5. Your weeks are Monday to Sunday. Sunday March 1 and Monday March 2 appear in the data. Which week does each fall in under resample('W-MON'), and what should you use?

    Show answer

    Under W-MON both land in the same bin that ends Monday March 2, which is not a Monday to Sunday week. Use resample('W-SUN'), whose bins run Monday to Sunday and are labelled by the Sunday. See time-series ops.

  6. An integer id column picks up one missing value after a merge, and the ids now print as 1042.0. Why, and how do you keep integers?

    Show answer

    NaN is a float, so pandas stores the whole column as floats, which prints and exports the ids with a decimal and loses exactness for very large ids. Use the nullable Int64 dtype, which keeps integers alongside a missing marker. See edge cases and traps.

Takeaway

  • State your approach, write the vectorized version first, and use apply only with a reason.
  • Always pass how and validate to merge, and check row counts before and after, because a non-unique key silently multiplies rows.
  • Every SQL window function has a pandas equivalent built from sort, groupby and cumsum, shift or rolling, and the sort order is part of the logic.
  • merge_asof is the point-in-time join: sort both sides, choose a direction, and set by and tolerance.
  • Time bucketing needs an explicit week anchor; W-SUN gives Monday to Sunday weeks.
  • Missing values change dtypes, so use nullable types when integer identity matters.
  • Knowing when to hand the work to SQL or DuckDB is part of the answer.

The next chapter applies these idioms to ten problems. → Pandas Problems

Further reading