Transform DataFrame and compute diff-in-diff
Company: Uber
Role: Data Scientist
Category: Data Manipulation (SQL/Python)
Difficulty: easy
Interview Round: Technical Screen
You are given a pandas DataFrame `df` with the following columns:
- `unit_id` (string): entity identifier (e.g., user, city, driver)
- `group` (string): either `'treatment'` or `'control'`
- `period` (string): either `'pre'` or `'post'`
- `y` (string): outcome stored as a string (should be numeric), with **exactly one missing value** (NaN)
Tasks:
1. Convert `y` from string to integer (assume all non-missing values are valid integer strings, e.g. `'12'`).
2. Impute the missing value in `y` using the **simple (unconditional) average** of the non-missing `y` values.
3. After steps (1)–(2), compute the **difference-in-differences (DiD)** estimate of the treatment effect on `y`:
\[
\text{DiD} = (\overline{y}_{\text{treat, post}} - \overline{y}_{\text{treat, pre}}) - (\overline{y}_{\text{ctrl, post}} - \overline{y}_{\text{ctrl, pre}})
\]
Return the scalar DiD estimate (and optionally the intermediate group-period means used).
Overview: This question evaluates proficiency in data cleaning (type conversion), missing-data handling, group-period aggregation, and estimating treatment effects via difference-in-differences.
Read the full Uber Data Scientist interview experience this question came from
Given table df(unit_id, group, period, y) where y is stored as a string and has exactly one missing value, cast y to integer, impute the missing y using the unconditional average of non-missing y values, then compute the Difference-in-Differences (DiD) estimate: (mean_y[treatment, post] - mean_y[treatment, pre]) - (mean_y[control, post] - mean_y[control, pre]). Return the scalar DiD (optionally include the intermediate group-period means).
Tables
df(unit_id VARCHAR, group VARCHAR, period VARCHAR, y VARCHAR)
Hints
- CAST y from VARCHAR to INTEGER first; use AVG() which ignores NULLs to get the unconditional mean.
- Impute with COALESCE(cast(y_int as decimal), global_avg). Bring the global average to every row via CROSS JOIN.
Community answers
Answer by SS
df['y'] = df['y'].astype(float) ## float can take care of nan values but not int
mean = df['y'].mean()
df['y'] = df['y'].fillna(mean)
y_treat_post = df[(df['group'] == 'treatment') & (df['period'] == 'post')]['y'].mean()
y_treat_pre = df[(df['group'] == 'treatment') & (df['period'] == 'pre')]['y'].mean()
y_control_post = df[(df['group'] == 'control') & (df['period'] == 'post')]['y'].mean()
y_control_pre = df[(df['group'] == 'control') & (df['period'] == 'pre')]['y'].mean()
did = (y_treat_post - y_treat_pre) - (y_control_post - y_control_pre)
Answer by sagarai9980
Question 1 .
df['y']=df['y'].astype(float)
Answer by sindhujakasula03
-- Write your SQL query herewith df1 as ( select avg(y::float) as avg_y from df),base_table as (select df.unit_id, df.group, df.period, coalesce(y::float, df1.avg_y) as y from df join df1 on 1=1),treatment_post as(select avg(y) as treat_post_mean from base_table where base_table.group = 'treatment' and base_table.period = 'post'),treatment_pre as(select avg(y) as treat_pre_mean from base_table where base_table.group = 'treatment' and base_table.period = 'pre'),control_pre as(select avg(y) as ctrl_pre_mean from base_table where base_table.group = 'control' and base_table.period = 'pre'),control_post as(select avg(y) as ctrl_post_mean from base_table where base_table.group = 'control' and base_table.period = 'post')select *, ((treat_post_mean-treat_pre_mean)-(ctrl_post_mean-ctrl_pre_mean)) as did_estimate from treatment_post join treatment_pre on 1=1 join control_pre on 1=1join control_post on 1=1