Learning Decision Trees by Inventing Them

1. Closer to 0. The bag is mostly one colour (80% Red), so it is fairly pure and easy to classify. Impurity should be low.
2. impurity(0.9, 0.1) is lower than impurity(0.8, 0.2). The 90/10 bag is more dominated by one colour, hence purer.
3. Because q = 1 - p, the function is really a function of one variable. It must equal 0 at the endpoints (p = 0 and p = 1) and reach its maximum at p = 0.5. It should be a smooth, symmetric, concave curve.
4. Yes. Swapping the label names (calling Red "Black" and Black "Red") does not change how mixed the bag is, so the function must be symmetric: impurity(p, q) == impurity(q, p).

Completed table:

The maximum is 0.25 at p = 0.5. The minimum is 0 at p = 0 and p = 1. The function has the right shape (zero at the extremes, peaks in the middle) but its peak is 0.25 instead of 1.

for p in [0.0, 0.2, 0.4, 0.5, 0.6, 0.8, 1.0]:
    print(f"{p:.1f}  {1-p:.1f}  {p*(1-p):.2f}")

plot_impurity_candidate(lambda p: p * (1 - p), title="f(p) = p * (1-p)")

Derivation:

• Maximum of f(p) = p*(1-p) is 0.25 at p = 0.5.
• Transformation: multiply by 4 to scale the peak from 0.25 to 1.
• Final formula: impurity(p, q) = 4 * p * q
• Check p=0.5: 4 * 0.5 * 0.5 = 1
• Check p=0: 4 * 0 * 1 = 0
• Check p=1: 4 * 1 * 0 = 0

def impurity(p, q):
    """
    Returns the impurity of a set with fraction p of Red and q of Black.
    impurity(0.5, 0.5) = 1  (most mixed)
    impurity(1, 0)     = 0  (all same)
    impurity(0, 1)     = 0
    """
    return 4 * p * q

# Verify:
print(impurity(0.5, 0.5))  # 1
print(impurity(0, 1))      # 0
print(impurity(1, 0))      # 0
print(impurity(0.8, 0.2))  # 0.64

# Ordering:
print(impurity(0.8, 0.2) < impurity(0.7, 0.3))  # True
print(impurity(0.7, 0.3) < impurity(0.6, 0.4))  # True

plot_impurity_candidate(lambda p: impurity(p, 1 - p),
                        title="My impurity function")

2. Bag A (impurity 0.36) is more certain than bag B (0.84). The impurity function agrees: a 90/10 split is purer than 70/30.
3. Yes, impurity(0.3, 0.7) == impurity(0.7, 0.3) == 0.84. This must hold because swapping colour labels does not change how mixed the bag is. Our formula 4*p*q is symmetric.
4. For 3 classes with equal fractions (1/3 each), you can extend to Gini-style: 1 - (p1^2 + p2^2 + p3^2) = 1 - 3*(1/3)^2 = 1 - 1/3 = 2/3 ≈ 0.667.

bags = [
    ("A", 0.9, 0.1),
    ("B", 0.7, 0.3),
    ("C", 0.6, 0.4),
    ("D", 0.5, 0.5),
    ("E", 0.3, 0.7),
    ("F", 1.0, 0.0),
]

for name, p, q in bags:
    print(f"Bag {name}: impurity({p}, {q}) = {impurity(p, q):.4f}")

# Results:
# Bag A: 0.3600   Bag B: 0.8400   Bag C: 0.9600
# Bag D: 1.0000   Bag E: 0.8400   Bag F: 0.0000

def list_impurity(items):
    """
    Computes the impurity of a list of "R" and "B" elements.
    Returns 0 for an empty list (no mixing possible).
    """
    if len(items) == 0:
        return 0
    total = len(items)
    count_r = items.count("R")
    count_b = total - count_r
    p = count_r / total
    q = count_b / total
    return impurity(p, q)

# Tests:
print(list_impurity(["R", "B", "B", "R", "R"]))  # 0.96
print(list_impurity(["R", "R", "R", "R"]))        # 0
print(list_impurity(["B", "B", "B", "B"]))        # 0
print(list_impurity(["R", "B"]))                  # 1.0
print(list_impurity([]))                          # 0

# Ranking:
list_A = ["R"]*8 + ["B"]*2  # p=0.8  imp=0.64
list_B = ["R"]*5 + ["B"]*5  # p=0.5  imp=1.0
list_C = ["R"]*7 + ["B"]*3  # p=0.7  imp=0.84
# Most mixed to least: B (1.0) > C (0.84) > A (0.64)

def total_impurity(left, right):
    """
    Returns the weighted total impurity of two groups.
    Each group's impurity is weighted by its size relative
    to the combined total.
    """
    n_left = len(left)
    n_right = len(right)
    n_total = n_left + n_right
    if n_total == 0:
        return 0
    return (n_left / n_total) * list_impurity(left) \
         + (n_right / n_total) * list_impurity(right)

# Tests:
print(total_impurity(["R","R","R"], ["B","B","B"]))
# 0.0 — perfectly separated, each group is pure

print(total_impurity(["R","B","R"], ["B","R","B"]))
# ~0.889 — both groups are 2/3 vs 1/3

print(total_impurity(["R"]*9 + ["B"], ["B"]))
# ~0.327 — left is slightly impure (10/11 weight),
#           right is pure (1/11 weight)

items = ["R", "B", "B", "R", "R"]

for k in range(1, len(items)):
    left = items[:k]
    right = items[k:]
    imp = total_impurity(left, right)
    print(f"k={k}: left={left}, right={right}, "
          f"total_impurity={imp:.4f}")

# k=0 gives left=[] (empty) — no meaningful split.
# k=5 gives right=[] (empty) — no meaningful split.
#
# Results:
# k=1: left=['R'],    right=['B','B','R','R'] → 0.75
# k=2: left=['R','B'],right=['B','R','R']     → 0.9333
# k=3: left=['R','B','B'], right=['R','R']    → 0.5333
# k=4: left=['R','B','B','R'], right=['R']    → 0.75
#
# k=3 gives the minimum total impurity (0.5333).
# Left=['R','B','B'] and Right=['R','R'] — the right
# group is pure and the left is mostly B.

def find_best_split_position(items):
    """
    Finds the split position k that minimises
    total_impurity(items[:k], items[k:]).
    Returns (best_k, min_impurity).
    """
    best_k = None
    min_imp = float('inf')
    for k in range(1, len(items)):
        left = items[:k]
        right = items[k:]
        imp = total_impurity(left, right)
        if imp < min_imp:
            min_imp = imp
            best_k = k
    return best_k, min_imp

# Tests:
for lst in [["R","B","B","R","R"],
            ["R","R","R","B","B","B"],
            ["B","R","B","R","B","R"],
            ["R","R","R","R","B"]]:
    best_k, min_imp = find_best_split_position(lst)
    print(f"{lst}")
    print(f"  best k={best_k}, left={lst[:best_k]}, "
          f"right={lst[best_k:]}, impurity={min_imp:.4f}")

# ["R","R","R","B","B","B"] → k=3 (obvious visual boundary)
# ["B","R","B","R","B","R"] → all positions are equally
#   mixed, so the first k with minimum is returned.

For ["R","R","R","B","B","B"] the minimum is sharp at k=3 because there is a clean boundary between the two colours.

For ["B","R","B","R","B","R"] the plot is nearly flat — no split position is much better than any other because the data is perfectly interleaved. There is no clear winner.

test_lists = [
    ["R", "B", "B", "R", "R"],
    ["R", "R", "R", "B", "B", "B"],
    ["B", "R", "B", "R", "B", "R"],
    ["R", "R", "R", "R", "B"],
]

for lst in test_lists:
    plot_split_scores(lst)

1. Around height 170-175, gender starts being predominantly Male. Below ~167 it is mostly Female.
2. A split line around 167-170 cm would separate M from F most cleanly.
3. A perfect separation is not possible because there are outliers: e.g. "M" at 168, "F" at 174, and "F" at 181. These points violate any single threshold.

The scatter plot shows F (red) clustered at the lower heights and M (blue) at the higher heights, with some overlap in the 168-181 range.

plot_height_gender(heights, genders)

The threshold t = 167 gives the lowest total impurity. At that split the left group (height <= 167) is [F, F, F, F] (pure) and the right group is mostly M with a few F outliers.

def split_by_threshold(heights, genders, t):
    """
    Splits the dataset at threshold t on height.
    Returns (left_genders, right_genders) where:
      left_genders  = genders where height <= t
      right_genders = genders where height > t
    """
    left_genders  = [g for h, g in zip(heights, genders) if h <= t]
    right_genders = [g for h, g in zip(heights, genders) if h > t]
    return left_genders, right_genders

# Test:
left, right = split_by_threshold(heights, genders, 172)
print("Left (height <= 172):", left)
print("Right (height > 172):", right)
print("Total impurity:", total_impurity(left, right))

# Scan thresholds:
for t in [160, 165, 170, 172, 175, 178, 182, 185, 188]:
    left, right = split_by_threshold(heights, genders, t)
    imp = total_impurity(left, right)
    print(f"t={t}: left={len(left)}, right={len(right)}, "
          f"impurity={imp:.4f}")

The best threshold is height <= 167. The left group is all F (pure). The right group is mostly M but contains a few F outliers (at 170, 174, 181), so the split is good but not perfect.

def find_best_threshold(heights, genders):
    """
    Finds the height threshold that minimises total_impurity.
    Returns (best_threshold, min_impurity).
    """
    best_t = None
    min_imp = float('inf')
    for t in sorted(set(heights)):
        left = [g for h, g in zip(heights, genders) if h <= t]
        right = [g for h, g in zip(heights, genders) if h > t]
        if not left or not right:
            continue
        imp = total_impurity(left, right)
        if imp < min_imp:
            min_imp = imp
            best_t = t
    return best_t, min_imp

best_t, best_imp = find_best_threshold(heights, genders)
print(f"Best threshold: height <= {best_t}")
print(f"Minimum total impurity: {best_imp:.4f}")

left, right = split_by_threshold(heights, genders, best_t)
print(f"Left group:  {left}")
print(f"Right group: {right}")

plot_height_gender(heights, genders, split_height=best_t)

The curve shows impurity dropping to a minimum at t = 167, then rising again. The minimum is relatively sharp, confirming that 167 is the best single-feature split for this dataset.

plot_threshold_scores(heights, genders)

1. salary=800, distance=20: salary <= 1000, so we go left → Reject (NO). We never even check distance.
2. salary=1200, distance=30: salary > 1000, go right → distance <= 40, go left → Accept (YES).
3. salary=1500, distance=60: salary > 1000, go right → distance > 40, go right → Reject (NO).
4. salary=1000, distance=any: The rule is salary > 1000. Since 1000 is NOT greater than 1000, we go left → Reject (NO). (Equivalently: salary <= 1000 means go left.)

class DecisionTreeNode:
    def __init__(self, decision=None, boundary=None,
                 boundary_value=None, left=None, right=None):
        self.decision = decision
        self.boundary = boundary
        self.boundary_value = boundary_value
        self.left = left
        self.right = right

    def check(self, features):
        if self.decision is not None:
            return self.decision
        if features[self.boundary] <= self.boundary_value:
            return self.left.check(features)
        else:
            return self.right.check(features)

def YES():
    return DecisionTreeNode(decision=True)

def NO():
    return DecisionTreeNode(decision=False)

# Build the job-offer tree:
job_tree = DecisionTreeNode(
    boundary="salary", boundary_value=1000,
    left=NO(),
    right=DecisionTreeNode(
        boundary="distance", boundary_value=40,
        left=YES(),
        right=NO()
    )
)

# Test:
print(job_tree.check({"salary": 800,  "distance": 20}))  # False
print(job_tree.check({"salary": 1200, "distance": 30}))  # True
print(job_tree.check({"salary": 1500, "distance": 60}))  # False
print(job_tree.check({"salary": 1000, "distance": 15}))  # False

Note: we split on coffee <= 0 (go left = NO) vs coffee > 0 (go right = YES). This is equivalent to checking coffee == 1.

job_tree_v2 = DecisionTreeNode(
    boundary="salary", boundary_value=1000,
    left=NO(),
    right=DecisionTreeNode(
        boundary="distance", boundary_value=40,
        left=DecisionTreeNode(
            boundary="coffee", boundary_value=0,
            left=NO(),     # coffee <= 0 → no coffee → reject
            right=YES()    # coffee > 0  → has coffee → accept
        ),
        right=NO()
    )
)

# Test:
tests = [
    {"salary": 1200, "distance": 30, "coffee": 1},  # True
    {"salary": 1200, "distance": 30, "coffee": 0},  # False
    {"salary": 1200, "distance": 60, "coffee": 1},  # False
    {"salary": 900,  "distance": 10, "coffee": 1},  # False
]
for job in tests:
    result = job_tree_v2.check(job)
    print(f"salary={job['salary']}, distance={job['distance']}, "
          f"coffee={job['coffee']}  →  {result}")

The dataset is loaded and printed. Notice that weight tends to be higher for M than for F, but there are some overlaps (e.g. F at weight 63 with height 174, and F at weight 66 with height 181). The question is: does splitting on weight produce a cleaner separation than splitting on height?

Weight typically wins as the best feature because it separates M and F more cleanly than height does (fewer outliers cross the weight boundary).

def find_decision_boundary(data, features, label):
    """
    Finds the best (feature, threshold) split across all
    given features.
    Returns (best_feature, best_threshold, min_impurity).
    """
    best_feature = None
    best_threshold = None
    min_imp = float('inf')

    for feature in features:
        unique_vals = sorted(set(row[feature] for row in data))
        for t in unique_vals:
            left  = [row[label] for row in data
                     if row[feature] <= t]
            right = [row[label] for row in data
                     if row[feature] > t]
            if not left or not right:
                continue
            imp = total_impurity(left, right)
            if imp < min_imp:
                min_imp = imp
                best_feature = feature
                best_threshold = t
    return best_feature, best_threshold, min_imp

best_feat, best_thresh, best_imp = find_decision_boundary(
    data, ["height", "weight"], "gender")
print(f"Best split: {best_feat} <= {best_thresh}")
print(f"Minimum total impurity: {best_imp:.4f}")

left_group = [row["gender"] for row in data
              if row[best_feat] <= best_thresh]
right_group = [row["gender"] for row in data
               if row[best_feat] > best_thresh]
print(f"Left:  {left_group}")
print(f"Right: {right_group}")

Comparing the two plots, weight gives a cleaner separation (lower impurity) because the M and F weight ranges overlap less than the height ranges.

best_h, best_ht, imp_h = find_decision_boundary(
    data, ["height"], "gender")
best_w, best_wt, imp_w = find_decision_boundary(
    data, ["weight"], "gender")

print(f"Best height split: height <= {best_ht}, "
      f"impurity = {imp_h:.4f}")
print(f"Best weight split: weight <= {best_wt}, "
      f"impurity = {imp_w:.4f}")

plot_feature_split(data, "height", "gender",
    split_value=best_ht,
    title=f"Height split at {best_ht} (imp={imp_h:.3f})")
plot_feature_split(data, "weight", "gender",
    split_value=best_wt,
    title=f"Weight split at {best_wt} (imp={imp_w:.3f})")

Manual trace using the height-only dataset:

1. Best first split: height <= 167 (found in Part 4).
2. Left group (height <= 167): genders = [F, F, F, F]. Right group (height > 167): genders = [M, F, M, M, F, M, M, M, M, M, F, M, M, M, M, M].
3. Left group is pure (all F). Right group is not pure (has 3 F among 16).
4. For the right group, the next best split isolates the remaining F outliers. The tree continues splitting until each leaf is pure.

height <= 167
├── Left: [F, F, F, F]  → LEAF: F (pure)
└── Right: [M, F, M, M, F, M, M, M, M, M, F, M, M, M, M, M]
    └── Split again on height (or weight if available)
        to isolate the remaining F outliers...

def majority_label(labels):
    """
    Returns the most common label in the list.
    Tie-breaking: returns whichever label appears first
    among those with the highest count.
    """
    if not labels:
        return None
    counts = {}
    for label in labels:
        counts[label] = counts.get(label, 0) + 1
    return max(counts, key=counts.get)

def is_pure(labels):
    """
    Returns True if all labels are the same (or the list
    is empty).
    """
    return len(set(labels)) <= 1

# Tests:
print(majority_label(["M", "M", "F", "M"]))  # M
print(majority_label(["F", "M"]))             # F (tie → first)
print(is_pure(["M", "M", "M"]))              # True
print(is_pure(["M", "M", "F"]))              # False
print(is_pure([]))                            # True

def build_tree(data, features, label):
    """
    Recursively builds a decision tree.
    Returns a DecisionTreeNode (leaf or boundary).
    """
    labels = [row[label] for row in data]

    # Base case 1: empty data
    if len(data) == 0:
        return DecisionTreeNode(decision=None)

    # Base case 2: pure or single element
    if is_pure(labels) or len(data) == 1:
        return DecisionTreeNode(decision=majority_label(labels))

    # Recursive case
    best_feat, best_thresh, best_imp = \
        find_decision_boundary(data, features, label)

    left_data  = [row for row in data
                  if row[best_feat] <= best_thresh]
    right_data = [row for row in data
                  if row[best_feat] >  best_thresh]

    # Safety: if split doesn't separate, return leaf
    if len(left_data) == 0 or len(right_data) == 0:
        return DecisionTreeNode(decision=majority_label(labels))

    left_tree  = build_tree(left_data,  features, label)
    right_tree = build_tree(right_data, features, label)

    return DecisionTreeNode(
        boundary=best_feat,
        boundary_value=best_thresh,
        left=left_tree,
        right=right_tree
    )

def print_tree(node, depth=0, label="ROOT"):
    indent = "  " * depth
    if node.decision is not None:
        print(f"{indent}[{label}] LEAF → {node.decision}")
    else:
        print(f"{indent}[{label}] {node.boundary} "
              f"<= {node.boundary_value}?")
        print_tree(node.left,  depth+1, "YES (left)")
        print_tree(node.right, depth+1, "NO (right)")

# Build and print:
tree = build_tree(data, ["height", "weight"], "gender")
print_tree(tree)

The tree achieves 100% accuracy on the training data because it keeps splitting until every leaf is pure. This is expected — with no depth limit the tree memorises the training set perfectly. Whether this generalises to new data is another question (see overfitting discussion in Part 9).

correct = 0
for row in data:
    predicted = tree.check(row)
    true_label = row["gender"]
    match = "Y" if predicted == true_label else "N"
    if predicted == true_label:
        correct += 1
    print(f"height={row['height']}, weight={row['weight']}: "
          f"true={true_label}, pred={predicted}  {match}")

print(f"\nAccuracy: {correct}/{len(data)} "
      f"= {correct/len(data)*100:.1f}%")

1. Features to split on: Extract them from the keys of the first dict in X: list(X[0].keys()). This way fit automatically discovers the available features.
2. predict returns: A list of predicted labels, one per row in X. Each prediction comes from traversing the tree for that row.
3. Leaf behaviour: A leaf stores the majority label (a string like "M" or "F"), and check() returns it directly.

class SimpleDecisionTree:
    def __init__(self):
        self.root = None
        self.features = None

    def fit(self, X, y):
        self.features = list(X[0].keys())
        label_col = "__label__"
        combined = [{**x, label_col: yi}
                    for x, yi in zip(X, y)]
        self.root = build_tree(combined,
                               self.features,
                               label_col)

    def predict(self, X):
        return [self.root.check(row) for row in X]

# Test:
X = [{"height": row["height"], "weight": row["weight"]}
     for row in data]
y = [row["gender"] for row in data]

clf = SimpleDecisionTree()
clf.fit(X, y)
predictions = clf.predict(X)

correct = sum(p == t for p, t in zip(predictions, y))
print(f"Training accuracy: {correct}/{len(y)} "
      f"= {correct/len(y)*100:.1f}%")

# Predict on new data:
new_people = [
    {"height": 163, "weight": 54},
    {"height": 185, "weight": 91},
    {"height": 172, "weight": 65},
]
new_preds = clf.predict(new_people)
for person, pred in zip(new_people, new_preds):
    print(f"  height={person['height']}, "
          f"weight={person['weight']}  →  {pred}")

The tree is several levels deep. Each split isolates one or more data points into pure leaves. The first split is on the feature (height or weight) that gives the best separation. Subsequent splits handle the remaining impure groups.

With more noise the tree would be deeper, because more exceptions force more splits to achieve purity. This is overfitting: the tree memorises the training data rather than learning a general pattern.

print_tree(clf.root)

With max_depth=1 the tree is a decision stump (one split). Training accuracy is lower but the model is simpler and less likely to overfit.

def build_tree_depth(data, features, label,
                     depth=0, max_depth=None):
    labels = [row[label] for row in data]

    if len(data) == 0:
        return DecisionTreeNode(decision=None)
    if is_pure(labels) or len(data) == 1:
        return DecisionTreeNode(decision=majority_label(labels))
    # NEW: stop if max depth reached
    if max_depth is not None and depth >= max_depth:
        return DecisionTreeNode(decision=majority_label(labels))

    best_feat, best_thresh, _ = \
        find_decision_boundary(data, features, label)
    left_data  = [r for r in data if r[best_feat] <= best_thresh]
    right_data = [r for r in data if r[best_feat] >  best_thresh]

    if not left_data or not right_data:
        return DecisionTreeNode(decision=majority_label(labels))

    return DecisionTreeNode(
        boundary=best_feat, boundary_value=best_thresh,
        left=build_tree_depth(left_data, features, label,
                              depth+1, max_depth),
        right=build_tree_depth(right_data, features, label,
                               depth+1, max_depth)
    )

# Test with max_depth = 1, 2, 3:
label_col = "__label__"
combined = [{**{"height": r["height"], "weight": r["weight"]},
             label_col: r["gender"]} for r in data]
feats = ["height", "weight"]
for md in [1, 2, 3, None]:
    t = build_tree_depth(combined, feats, label_col,
                         max_depth=md)
    preds = [t.check(r) for r in data]
    acc = sum(p == r["gender"] for p, r in zip(preds, data))
    print(f"max_depth={md}: accuracy={acc}/{len(data)}")
    print_tree(t)

Higher min_samples produces a shallower tree with fewer splits, trading training accuracy for simplicity and better generalisation.

def build_tree_ms(data, features, label, min_samples=1):
    labels = [row[label] for row in data]

    if len(data) == 0:
        return DecisionTreeNode(decision=None)
    if is_pure(labels) or len(data) == 1:
        return DecisionTreeNode(decision=majority_label(labels))
    # NEW: stop if too few samples
    if len(data) < min_samples:
        return DecisionTreeNode(decision=majority_label(labels))

    best_feat, best_thresh, _ = \
        find_decision_boundary(data, features, label)
    left_data  = [r for r in data if r[best_feat] <= best_thresh]
    right_data = [r for r in data if r[best_feat] >  best_thresh]

    if not left_data or not right_data:
        return DecisionTreeNode(decision=majority_label(labels))

    return DecisionTreeNode(
        boundary=best_feat, boundary_value=best_thresh,
        left=build_tree_ms(left_data, features, label,
                           min_samples),
        right=build_tree_ms(right_data, features, label,
                            min_samples)
    )

# Compare min_samples = 1, 3, 5:
for ms in [1, 3, 5]:
    t = build_tree_ms(combined, feats, label_col,
                      min_samples=ms)
    preds = [t.check(r) for r in data]
    acc = sum(p == r["gender"] for p, r in zip(preds, data))
    print(f"min_samples={ms}: accuracy={acc}/{len(data)}")
    print_tree(t)
    print()

The tree should achieve 100% training accuracy. The first split is usually on petal width (cm) at around 0.8 cm, perfectly separating setosa from the other two species. Test accuracy is typically 93-97%.

from sklearn.datasets import load_iris

iris = load_iris()
feature_names = iris.feature_names
X_iris = [dict(zip(feature_names, row))
          for row in iris.data]
y_iris = [iris.target_names[t] for t in iris.target]

clf_iris = SimpleDecisionTree()
clf_iris.fit(X_iris, y_iris)

preds = clf_iris.predict(X_iris)
acc = sum(p == t for p, t in zip(preds, y_iris))
print(f"Iris training accuracy: "
      f"{acc}/{len(y_iris)} = {acc/len(y_iris)*100:.1f}%")

print("\nTree structure:")
print_tree(clf_iris.root)

# Bonus: train/test split
import random
random.seed(42)
indices = list(range(len(X_iris)))
random.shuffle(indices)
split = int(0.8 * len(indices))
train_idx, test_idx = indices[:split], indices[split:]

X_train = [X_iris[i] for i in train_idx]
y_train = [y_iris[i] for i in train_idx]
X_test  = [X_iris[i] for i in test_idx]
y_test  = [y_iris[i] for i in test_idx]

clf2 = SimpleDecisionTree()
clf2.fit(X_train, y_train)
test_preds = clf2.predict(X_test)
test_acc = sum(p == t for p, t in zip(test_preds, y_test))
print(f"\nTest accuracy: {test_acc}/{len(y_test)} "
      f"= {test_acc/len(y_test)*100:.1f}%")

For two classes: Gini = 1 - p^2 - q^2 = 2pq. Your formula 4pq differs from Gini by a factor of 2, but this does not affect which split wins — both rank all possible splits identically. Entropy has a slightly different shape (it peaks at 1.0 instead of 0.5) but also ranks splits the same way.

def gini_impurity(p, q):
    """
    Gini impurity for two classes.
    Gini = 1 - (p^2 + q^2) = 2*p*q
    Ranges from 0 (pure) to 0.5 (maximally mixed).
    """
    return 1 - (p**2 + q**2)

def entropy_impurity(p, q):
    """
    Entropy for two classes.
    H = -(p*log2(p) + q*log2(q))
    Handle p=0 or q=0 as contributing 0.
    Ranges from 0 (pure) to 1.0 (maximally mixed).
    """
    result = 0
    if p > 0:
        result -= p * math.log2(p)
    if q > 0:
        result -= q * math.log2(q)
    return result

# Plot all three:
ps = [i/100 for i in range(101)]
plt.figure(figsize=(8, 4))
plt.plot(ps, [impurity(p, 1-p) for p in ps],
         label='Your formula (4pq)')
plt.plot(ps, [gini_impurity(p, 1-p) for p in ps],
         label='Gini (2pq)', linestyle='--')
plt.plot(ps, [entropy_impurity(p, 1-p) for p in ps],
         label='Entropy', linestyle=':')
plt.xlabel('p'); plt.ylabel('impurity')
plt.title('Comparing impurity measures')
plt.legend(); plt.grid(True); plt.show()