Heterogeneous Effect

Problem

Estimate the treatment effect in two subgroups split on X at threshold: for each side compute mean(y | treatment=1) - mean(y | treatment=0).

A side without both treated and control samples gets 0.0. Return [effect_left, effect_right].

Hint

Comparing treated and control means within each subgroup is a simple CATE estimator.

case 1

in: [[1,1,1,1,9,9,9,9],[3,1,4,2,9,1,7,3],[1,0,1,0,1,0,1,0],5]

expected: [2,6]

case 2

in: [[1,2,3,4],[5,6,7,8],[1,1,0,0],2.5]

expected: [0,0]

case 3

in: [[1,2,3,4,5,6],[1,2,3,4,5,6],[1,0,1,0,1,0],3.5]

expected: [0,0]

Starter codePython

def heterogeneous_effect(X, y, treatment, threshold):
    # Returns [effect_left, effect_right]
    # Your code here
    pass

3 test cases · run your code against them in the editor