Confusion Matrix

Implement a multiclass confusion matrix from scratch with NumPy, without calling a library metric. The confusion matrix is the foundation for accuracy, precision, recall, and F1, and interviewers use it to check that you can accumulate counts into a fixed-size grid without a slow Python loop.

Definition

Given n samples with integer true labels y_true and predicted labels y_pred, both in [0, num_classes), the confusion matrix M is a (k, k) array where

M[i, j] = number of samples whose true class is i and whose predicted class is j

Rows index the true class, columns index the predicted class. The diagonal M[i, i] counts correct predictions for class i; off-diagonal entries are the specific mistakes.

Task

Complete confusion_matrix(y_true, y_pred, num_classes) so it returns a (k, k) integer NumPy array following the definition above, with k = num_classes. Classes that never appear must still get an all-zero row and column, so the shape depends only on num_classes, not on which labels are present. Do not call sklearn or any built-in confusion-matrix helper.

Example

y_true = np.array([0, 1, 2, 2, 0, 1])
y_pred = np.array([0, 2, 2, 0, 0, 1])
confusion_matrix(y_true, y_pred, num_classes=3)
# -> array([[2, 0, 0],
#           [0, 1, 1],
#           [1, 0, 1]])

Class 0 is predicted correctly both times (M[0, 0] = 2). One class-1 sample is misread as class 2 (M[1, 2] = 1), and one class-2 sample is misread as class 0 (M[2, 0] = 1).

Constraints

• 1 <= n <= 10^6 and 1 <= num_classes <= 10^3; use vectorized NumPy, not a Python loop over samples.
• All labels are valid integers in [0, num_classes).

Approach

Allocate a (k, k) zero matrix, then scatter one +1 into cell [y_true[i], y_pred[i]] for every sample. np.add.at does exactly this unbuffered scatter-add, so repeated (true, pred) pairs accumulate correctly. An equivalent trick is to flatten each pair to a single index true * k + pred and count with np.bincount, then reshape to (k, k).

Solution

import numpy as np

def confusion_matrix(y_true, y_pred, num_classes):
    k = num_classes
    M = np.zeros((k, k), dtype=int)
    np.add.at(M, (y_true, y_pred), 1)   # scatter-add +1 per (true, pred) pair
    return M

Bincount variant, same result and often a bit faster:

def confusion_matrix(y_true, y_pred, num_classes):
    k = num_classes
    flat = np.asarray(y_true) * k + np.asarray(y_pred)
    counts = np.bincount(flat, minlength=k * k)
    return counts.reshape(k, k)

Walkthrough

On the example y_true = [0, 1, 2, 2, 0, 1], y_pred = [0, 2, 2, 0, 0, 1] with k = 3:

• Pairs (true, pred) are (0,0), (1,2), (2,2), (2,0), (0,0), (1,1).
• np.add.at adds +1 at each of those cells. (0,0) occurs twice, so M[0,0] = 2.
• The remaining pairs each fire once: M[1,2] = 1, M[2,2] = 1, M[2,0] = 1, M[1,1] = 1.
• Everything else stays zero, giving [[2,0,0],[0,1,1],[1,0,1]].

In the bincount variant those same pairs flatten to indices [0, 5, 8, 6, 0, 4] (using true*3+pred); np.bincount with minlength=9 counts them into [2,0,0,0,1,1,1,0,1], which reshapes to the identical matrix.

Complexity & notes

• Time O(n + k^2), space O(k^2) — one scatter over n samples plus allocating and returning the k x k grid.
• Use np.add.at, not M[y_true, y_pred] += 1. Fancy-index assignment buffers, so duplicate pairs would overwrite rather than accumulate and you would undercount; np.add.at is the unbuffered version that adds every occurrence.
• Fixing the shape at (num_classes, num_classes) guarantees absent classes keep their zero row and column, which matters when you later slice per-class precision or recall.
• From here, per-class recall is M[i, i] / M[i, :].sum(), per-class precision is M[i, i] / M[:, i].sum(), and overall accuracy is np.trace(M) / M.sum().
• If labels could be out of range, validate or clip them first; the scatter would raise an IndexError, and the bincount variant would silently land in the wrong cell.