Decision trees with Gini splits and naive Bayes from scratch, step by step

In this article

1. Before you start
2. Why interviewers ask this
3. Decision trees and naive Bayes from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one step of the machine learning (ML) coding round learning path. You write plain Python here, with no libraries.

Why interviewers ask this

Candidates report being asked to write a decision tree with Gini impurity, or a classical model such as bagging or k-means, in an ML coding round. The reason is practical: a from-scratch tree or naive Bayes needs no framework, so the interviewer sees your loops, your counting and your handling of edge cases directly.

A weak answer recites "the tree picks the best split." A strong answer names the impurity formula, tries every midpoint between sorted values, weights each child by its size, and stops on a depth or leaf-size limit. For naive Bayes, a strong answer smooths the counts and sums log probabilities. A Staff-level answer also says what the model gets wrong: a greedy tree is not certain to be the best tree, and naive Bayes probabilities are poor even when its class choices are good.

Decision trees and naive Bayes from first principles

Gini impurity

Gini impurity is a number that measures how mixed the classes in one group of rows are: one minus the sum of the squared class shares. It is zero when every row has the same class and grows as the classes mix more evenly.

Split threshold

A split threshold is a cut value on one numeric feature that sends rows with a value at or below it to the left child and the rest to the right child. A common choice is the midpoint between two neighboring sorted values.

Log score

A log score is the sum of the logarithm of a class prior and the logarithms of its feature likelihoods. It ranks classes the same way as the raw product but cannot underflow to zero.

In plain words

Think of the game of twenty questions. Each question splits the possible answers, and a good question leaves two piles that are each mostly one kind of thing. A decision tree plays that game with your data. The analogy breaks in one place: the tree must pick each question from numbers alone, so it needs a score for how clean a pile is.

Naive Bayes plays a different game. It asks, for each class, how well the words in a message match what that class usually contains. Each word votes independently, and the votes are combined.

In one precise sentence, a tree learns thresholds that purify groups, and naive Bayes learns per-class word frequencies and combines them under a simplifying assumption.

Step 1: Gini impurity scores how mixed one node is

Gini impurity of a node is one minus the sum of the squared class shares. For the toy data below (ten rows, class 0 on four rows and class 1 on six; illustrative), the shares are 0.4 and 0.6, so the impurity is 1 − (0.16 + 0.36) = 0.48.¹

Input: the labels of the rows in one node. Output: one number between zero and one half for two classes. A node holding a single class scores zero, which tells the tree to stop. Entropy, minus the sum of each share times its base-2 logarithm, is an alternative score; on this toy data it is 0.971 at the root.¹ You now have a way to say how mixed a group is, which the next step uses to compare cuts.

Step 2: Pick the threshold with the lowest weighted Gini

The best split is the threshold whose two children have the lowest impurity averaged by size. Sort the distinct values of a feature, take the midpoint of each neighboring pair as a candidate threshold, and score each one. The scoring formula is the left share times the left Gini plus the right share times the right Gini.

The common error is a plain average of the two child impurities. On the toy data below, the plain average is lowest at the threshold 6 (0.1224), so it picks 6; weighting by size picks 9, with 0.16 against 0.1714 at 6.¹ Weighting stops a small pure child from hiding a large mixed one.

On the toy data, the threshold 9 sends five rows left (four of class 0, one of class 1, Gini 0.32) and five right (all class 1, Gini 0). The weighted Gini is 0.5 × 0.32 + 0.5 × 0 = 0.16, a drop of 0.32 from the root.¹ Input: node rows. Output: one threshold. Entropy also picks 9 on this data.¹

Step 3: Grow the tree and know when to stop

A decision tree is grown by repeating Step 2 on each child until a stopping rule fires. Typical rules are a maximum depth, a minimum number of rows in a leaf, or a node with only one class. A leaf predicts the most common class among its rows.

This greedy method chooses the best cut at each node and never revisits it. The scikit-learn decision-tree guide states that such locally optimal decisions cannot ensure the globally optimal tree, and that an ensemble of trees helps. The same page names pruning, a minimum leaf size and a maximum depth as ways to stop overfitting, which is when a model memorizes training rows instead of learning the pattern.

On the toy data, a depth-one tree has two leaves and gets nine of ten training rows right (0.9).¹ An unlimited tree keeps splitting the left child until one leaf holds the single row x = 7, and fits all ten, a fit that is probably memorization.¹

Step 4: Naive Bayes counts words per class, with smoothing

Naive Bayes is a classifier that picks the class with the highest prior times the product of its word likelihoods. The "naive" part is conditional independence: the assumption that, once the class is known, the features do not influence each other. The scikit-learn naive Bayes guide describes the method as Bayes' theorem with exactly this assumption between every pair of features.

Training is counting. The toy corpus (illustrative) has five short messages: two spam and three ham, with a vocabulary of eleven words. The priors are 0.4 for spam and 0.6 for ham. Spam messages contain six words in total and ham messages nine.¹

A word's likelihood is its count plus one, divided by the class's total words plus the vocabulary size. The same page defines the smoothing parameter alpha, and calls alpha equal to one Laplace smoothing. For the word "free" in spam that is (1 + 1) / (6 + 11) = 0.1176.¹

Step 5: Add logs instead of multiplying, then pick the class

Naive Bayes adds the logarithms of the prior and of each word's likelihood, then picks the class with the larger sum. A product of many small probabilities falls below the smallest number a float can hold, so it becomes zero; a sum of logs stays finite. For a thousand probabilities of 0.01 (illustrative), the product evaluates to 0.0 in Python while the sum of logs is -4605.17.¹ Manning, Raghavan and Schütze (2009) likewise compute naive Bayes in log space to avoid floating-point underflow. The same log-space idea underlies a numerically stable softmax.

For the message "free money lunch", the spam log score is -7.6242 and the ham log score is -7.7063, so the model picks spam.¹ Input: word counts and a message. Output: a class. Turn the scores back into probabilities and spam gets 0.5205, ham 0.4795; treat those as ranks, not facts, because the same documentation calls naive Bayes a decent classifier but a bad probability estimator.¹

1 · Gini impurity2 · Best threshold3 · Grow and stop4 · Add-one counts5 · Log scoreReplay

Step-by-step description

1. Score how mixed one node is.
2. Keep the lowest weighted Gini.
3. Recurse until a stopping rule fires.
4. Count words per class, add one.
5. Add logs and pick the class.

• The best split of the toy data is the threshold 9 with a weighted Gini of 0.16.
• Spam scores -7.6242 and ham scores -7.7063, so the model picks spam.

Worked example

The tree data is ten rows with one numeric feature x and a class (illustrative). The x values are 2, 3, 5, 7, 8, 10, 11, 13, 14, 16 and the classes are 0, 0, 0, 1, 0, 1, 1, 1, 1, 1.¹ The candidate thresholds are the nine midpoints of neighboring values.

The table lists five of the nine thresholds; the chart below shows eight of them, adding 4 (0.3) and 15 (0.4444) and leaving out 13.5 (0.4).¹ The weighted Gini by threshold is lowest at the threshold 9, where it is 0.16. A second dip at the threshold 6 reaches 0.1714.¹

Text description of the diagram

• The weighted Gini by threshold is lowest at the threshold 9, where it is 0.16.
• A second dip at the threshold 6 reaches 0.1714.

For naive Bayes, the message "free money lunch" is scored against the corpus from Step 4.¹

Without smoothing, "free" never appears in ham, so its ham likelihood is zero and the whole ham product collapses to zero whatever the other words say.¹

The tree in Python

A from-scratch tree needs four functions: gini, best_split, build_tree and predict_tree. The code below is adapted from the saved script. gini scores a list of labels. best_split tries the midpoint between each pair of neighboring sorted values, scores each cut by weighted Gini, and skips any cut that leaves a child smaller than min_leaf, which also skips empty children. build_tree recurses until the node is pure, max_depth is reached or no split is allowed, and a leaf stores the majority class. predict_tree walks from the root to a leaf.

from collections import Counter

def gini(ys):
    n = len(ys)
    if n == 0:
        return 0.0
    return 1.0 - sum((c / n) ** 2 for c in Counter(ys).values())

def best_split(xs, ys, min_leaf=1):
    best = None  # (weighted Gini, threshold)
    vals = sorted(set(xs))
    for a, b in zip(vals, vals[1:]):
        t = (a + b) / 2                      # midpoint candidate
        left = [y for x, y in zip(xs, ys) if x <= t]
        right = [y for x, y in zip(xs, ys) if x > t]
        if len(left) < min_leaf or len(right) < min_leaf:
            continue                         # skip tiny or empty children
        w = (len(left) * gini(left) + len(right) * gini(right)) / len(ys)
        if best is None or w < best[0]:      # strict <: ties keep the first
            best = (w, t)
    return best

def build_tree(xs, ys, depth=0, max_depth=3, min_leaf=1):
    majority = Counter(ys).most_common(1)[0][0]
    if depth == max_depth or gini(ys) == 0:
        return majority                      # leaf
    split = best_split(xs, ys, min_leaf)
    if split is None:
        return majority
    t = split[1]
    left = [(x, y) for x, y in zip(xs, ys) if x <= t]
    right = [(x, y) for x, y in zip(xs, ys) if x > t]
    return {"t": t,
            "left": build_tree([x for x, _ in left], [y for _, y in left], depth + 1, max_depth, min_leaf),
            "right": build_tree([x for x, _ in right], [y for _, y in right], depth + 1, max_depth, min_leaf)}

def predict_tree(node, x):
    while isinstance(node, dict):
        node = node["left"] if x <= node["t"] else node["right"]
    return node

On the ten toy rows with max_depth=1, the code returns the root threshold 9, a left leaf 0 and a right leaf 1, and gets 0.9 of the training rows right.¹ The version above rescans all rows for every threshold, so scoring one feature at one node costs O(n²) for n rows. Sorting once and keeping running class counts as the cut moves costs O(n log n), dominated by the sort.

Naive Bayes in Python

Naive Bayes needs two short functions. The first, train_nb, only counts: how many messages each class has, how many times each word appears per class, and the vocabulary. The second, predict_nb, starts each class at the log of its prior and adds the log of (count + 1) / (class total + vocabulary size) for every word of the message, then returns the class with the larger sum.

import math
from collections import Counter

def train_nb(docs):                          # docs: list of (text, label)
    classes = Counter(label for _, label in docs)
    counts = {c: Counter() for c in classes}
    vocab = set()
    for text, label in docs:
        words = text.split()
        counts[label].update(words)
        vocab.update(words)
    return classes, counts, vocab, len(docs)

def predict_nb(model, text):
    classes, counts, vocab, n = model
    scores = {}
    for c in classes:
        total = sum(counts[c].values())
        s = math.log(classes[c] / n)         # log prior
        for w in text.split():
            s += math.log((counts[c][w] + 1) / (total + len(vocab)))
        scores[c] = s
    return max(scores, key=scores.get), scores

On the five-message corpus, the message "free money lunch" returns spam, with log scores -7.6242 for spam and -7.7063 for ham.¹ Training is one pass over the data and prediction is one pass over the message words per class.

Warm-up: mean, median and mode in core Python

Candidates report this warm-up paired with the tree question, so know it cold. Sort the list for the median, and count occurrences for the mode, keeping every value tied for the top count.

from collections import Counter

def mean(xs):
    return sum(xs) / len(xs)

def median(xs):
    s = sorted(xs)
    n = len(s)
    mid = n // 2
    return s[mid] if n % 2 else (s[mid - 1] + s[mid]) / 2

def mode(xs):
    counts = Counter(xs)
    top = max(counts.values())
    return sorted(k for k, v in counts.items() if v == top)

On the list 4, 1, 3, 3, 9, 7, 3, 1 (illustrative), the mean is 3.875, the median is 3.0 and the mode is 3.¹ State what you do for an empty list; the functions above raise an error, which is a reasonable choice if you say so. For a median over several factors, sort by a key tuple such as the main score and then a tie-breaker, and apply the same middle rule; for a weighted median, walk the sorted values until the running weight reaches half of the total.

Trade-offs and edge cases

A decision tree needs no assumption about feature independence and can capture interactions, but it overfits when grown deep and is unstable. The scikit-learn decision-tree guide notes that small variations in the data can produce a very different tree. Bagging, which trains many trees on resampled rows and takes a vote, is the standard remedy.

Naive Bayes trains in one pass over counts and works with few rows, but its independence assumption is usually false and its probabilities are unreliable. Smoothing is needed for any word missing from a class; an alpha below one is called Lidstone smoothing by the same documentation.

Text description of the diagram

• A decision tree overfits when it is grown deep.
• Naive Bayes needs smoothing so a zero count does not remove a class.

Edge cases for the tree code: a feature with one distinct value yields no candidate thresholds, so return a leaf; ties between thresholds need a fixed rule, such as keeping the first; a split that leaves an empty child must be skipped; and a minimum leaf size should block splits that create tiny leaves. To judge the toy tree beyond its 0.9 training accuracy, see evaluation metrics from scratch.¹

How to answer it in an interview

The spoken structure: say the split score, say the stopping rule and the greedy limit, then say the counting, smoothing and log-space steps.

Common mistakes

• "I average the two child Gini values." (wrong) A tiny child counts as much as a big one. (right) "I weight each child's Gini by its share of the rows."
• "The tree finds the best possible tree." (wrong) It picks the best cut at each node only. (right) "It is greedy, so I use depth limits and an ensemble."
• "Unseen words get probability zero." (wrong) One missing word then wipes out a class. (right) "I add one to every count so no likelihood is zero."
• "I multiply the probabilities." (wrong) Long messages underflow to zero. (right) "I sum log probabilities and compare the sums."

Follow-up questions

How do you implement a decision tree from scratch in Python with Gini impurity?

Write a gini function over labels, a best_split that scores every midpoint by weighted Gini, and a recursive build_tree that stops on purity, depth or leaf size. Prediction walks from the root to a leaf by comparing the feature with each threshold.

How do you choose the threshold on a continuous feature?

Sort the distinct values and test the midpoint between each neighboring pair, because only those cuts change which rows go left. The best one has the lowest weighted impurity; the toy data above picks 9.

Why does naive Bayes work when features are not independent?

Classification only needs the correct class to have the highest score, not exact probabilities. The scikit-learn documentation calls naive Bayes a decent classifier but a bad estimator, so use the ranking and distrust the percentages.

Is Gini or entropy better for choosing splits?

On the toy data both pick the same threshold, 9; they may agree often, but check on your own data, because this page has no source for a general claim. Gini avoids a logarithm and is a little cheaper to compute.

How would you implement bagging on top of your tree?

Draw several resamples of the rows with replacement, train one tree on each, and predict by majority vote. This reduces the instability of a single greedy tree; the same from-scratch discipline applies to scaled dot-product attention and to the forward and backward pass of a small network.

Sources

1. Decision Trees - scikit-learn user guide — scikit-learn 1.x stable documentation (accessed 2026-10-26), 26 Oct 2026
2. Naive Bayes - scikit-learn user guide — scikit-learn 1.x stable documentation (accessed 2026-10-26), 26 Oct 2026
3. Naive Bayes text classification - Introduction to Information Retrieval — Cambridge University Press (Manning, Raghavan, Schütze), 7 Apr 2009