One-Hot Encoding

Convert a 1-D array of integer class labels into a one-hot encoded matrix. One-hot encoding is the standard way to feed categorical labels into models that expect a vector target (softmax classifiers, cross-entropy loss), and interviewers use it to check that you can turn a per-row index into a matrix with NumPy indexing instead of a loop.

Definition

Given n labels and num_classes columns, the output is an (n, num_classes) matrix where row i is all zeros except for a single 1 in column labels[i].

Task

Complete one_hot_encode(labels, num_classes) so it returns the (n, num_classes) one-hot matrix. labels is a 1-D NumPy array of integers, each in [0, num_classes). Use vectorized NumPy indexing; do not write a Python loop and do not call sklearn or any built-in one-hot helper.

Example

labels = np.array([0, 2, 1, 2])
one_hot_encode(labels, 3)
# -> array([[1, 0, 0],
#           [0, 0, 1],
#           [0, 1, 0],
#           [0, 0, 1]])

Row 0 has its 1 in column 0, row 1 in column 2, row 2 in column 1, and row 3 in column 2.

Constraints

• 1 <= n <= 10^6 and 1 <= num_classes <= 10^4; use vectorized NumPy, not a Python loop.
• Every value in labels is a valid class index in [0, num_classes).

Approach

Allocate an (n, num_classes) matrix of zeros, then use fancy indexing to drop a single 1 into each row at the column named by its label. Pairing np.arange(n) (the row indices) with labels (the column indices) addresses exactly one cell per row, so a single assignment fills the whole matrix without a loop.

Solution

import numpy as np

def one_hot_encode(labels: np.ndarray, num_classes: int) -> np.ndarray:
    n = labels.shape[0]
    out = np.zeros((n, num_classes), dtype=int)
    out[np.arange(n), labels] = 1        # place one 1 per row via fancy indexing
    return out

Walkthrough

On the example labels = [0, 2, 1, 2], num_classes = 3:

• n = 4, so out starts as a 4 x 3 block of zeros.
• np.arange(n) = [0, 1, 2, 3] are the row indices, paired elementwise with labels = [0, 2, 1, 2] as column indices.
• The assignment sets out[0,0], out[1,2], out[2,1], out[3,2] to 1 simultaneously.
• Result: [[1,0,0], [0,0,1], [0,1,0], [0,0,1]].

Complexity & notes

• Time O(n * num_classes) to allocate and zero the output matrix, plus O(n) for the assignment; space O(n * num_classes) for the result itself.
• An equivalent one-liner is np.eye(num_classes, dtype=int)[labels], which indexes the rows of the identity matrix; it is concise but builds a num_classes x num_classes identity first, so it wastes memory when num_classes is large.
• Fancy indexing is the key idea: passing two integer arrays of equal length to out[rows, cols] addresses the paired cells, not a cross-product.
• If labels could fall outside [0, num_classes) you would validate first; NumPy would raise an IndexError on out-of-range column indices, and negative values would silently wrap from the end.