Skip to content

Instantly share code, notes, and snippets.

@eugeneyan
Last active February 21, 2021 19:06
Show Gist options
  • Save eugeneyan/966990b869cdaed33fbbb2bbd6a6644f to your computer and use it in GitHub Desktop.
Save eugeneyan/966990b869cdaed33fbbb2bbd6a6644f to your computer and use it in GitHub Desktop.
Test Gini impurity and gain
def test_gini_impurity():
assert round(gini_impurity([1, 1, 1, 1, 1, 1, 1, 1]), 3) == 0
assert round(gini_impurity([1, 1, 1, 1, 1, 1, 1, 0]), 3) == 0.219
assert round(gini_impurity([1, 1, 1, 1, 1, 1, 0, 0]), 3) == 0.375
assert round(gini_impurity([1, 1, 1, 1, 1, 0, 0, 0]), 3) == 0.469
assert round(gini_impurity([1, 1, 1, 1, 0, 0, 0, 0]), 3) == 0.500
def test_gini_gain():
assert round(gini_gain([1, 1, 1, 1, 0, 0, 0, 0], [[1, 1, 1, 1], [0, 0, 0, 0]]), 3) == 0.5
assert round(gini_gain([1, 1, 1, 1, 0, 0, 0, 0], [[1, 1, 1, 0], [0, 0, 0, 1]]), 3) == 0.125
assert round(gini_gain([1, 1, 1, 1, 0, 0, 0, 0], [[1, 0, 0, 0], [0, 1, 1, 1]]), 3) == 0.125
assert round(gini_gain([1, 1, 1, 1, 0, 0, 0, 0], [[1, 1, 0, 0], [0, 0, 1, 1]]), 3) == 0.0
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment