0

One-hot encode an integer array

by Frontier Interview· 1 day ago· ⑂ 0 remixes

Question: One-hot encode an integer array

Companies: Netflix

Difficulty: Easy

Problem statement

Write one_hot(arr: list[int]) -> list[list[int]] that returns the one-hot encoding of an integer array as a 2D list.

  • The matrix has one row per element of arr, in the same order.
  • It has one column per distinct value in arr. Columns are ordered by value, smallest first.
  • Row i has a 1 in the column for arr[i] and 0 everywhere else.

An empty input returns an empty list.

Constraints

  • 0 <= len(arr) <= 10^5
  • -10^9 <= arr[i] <= 10^9
  • The number of distinct values times len(arr) fits in memory.

Example 1

arr = [4, 1, 4, 7]
columns (sorted distinct values): [1, 4, 7]
output:
[[0, 1, 0],
 [1, 0, 0],
 [0, 1, 0],
 [0, 0, 1]]

Example 2

arr = [3, 3, -2]
columns: [-2, 3]
output: [[0, 1], [0, 1], [1, 0]]

Possible next parts: decode a matrix back to the values, return a compact representation when there are many distinct values, and encode new data with columns fixed from earlier data.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

The columns are the distinct values in ascending order, so the work splits into two steps:

  1. Build the column index. Take sorted(set(arr)). That list is the header: column j stands for the j-th smallest distinct value. Store it as a dict value -> column so each lookup is O(1).
  2. Fill the rows. For every element, make a row of k zeros (k = number of distinct values) and set the cell at index[x] to 1.

Brute force first. A first draft often does columns.index(x) for every element. That is a linear scan per row, O(n·k) lookups on top of the O(n·k) output. It is correct, but the dict makes the cost clear: lookups become O(1), and building the matrix is the only O(n·k) part. You can't avoid that part, because it is the size of the output.

Invariants.

  • Every row has exactly one 1, and it sits in the column whose value equals arr[i].
  • The column order depends only on the set of values, not on where they appear. Equal elements get identical rows.
  • Every column has at least one 1, because every distinct value appears at least once.

Sparse view. When k is large (for example IDs), the dense matrix is mostly zeros. The useful representation is then the pair (columns, codes), where codes[i] = index[arr[i]]. The dense matrix can be rebuilt from it whenever it's needed. The code includes encode_sparse and decode to show the round trip, since interviewers often ask for this next.

Code

from typing import Hashable, Iterable, Sequence


def one_hot(arr: Sequence[int]) -> list[list[int]]:
    """Dense one-hot matrix: one row per element, one column per distinct
    value in ascending order."""
    columns = sorted(set(arr))
    index = {value: j for j, value in enumerate(columns)}
    k = len(columns)
    matrix: list[list[int]] = []
    for x in arr:
        row = [0] * k
        row[index[x]] = 1
        matrix.append(row)
    return matrix


def encode_sparse(arr: Sequence[Hashable]) -> tuple[list, list[int]]:
    """Return (columns, codes) where codes[i] is the column of arr[i].
    O(n) memory instead of O(n * k)."""
    columns = sorted(set(arr))
    index = {value: j for j, value in enumerate(columns)}
    return columns, [index[x] for x in arr]


def decode(columns: Sequence, matrix: Iterable[Sequence[int]]) -> list:
    """Invert a dense one-hot matrix back to the original values."""
    out = []
    for r, row in enumerate(matrix):
        hot = [j for j, bit in enumerate(row) if bit == 1]
        if len(hot) != 1 or sum(row) != 1:
            raise ValueError(f"row {r} is not one-hot: {list(row)}")
        out.append(columns[hot[0]])
    return out


class OneHotEncoder:
    """Fit on training data, then transform new data with the same columns.
    Unseen values become an all-zero row (or raise, if strict)."""

    def __init__(self, strict: bool = False) -> None:
        self.strict = strict
        self.columns: list = []
        self._index: dict = {}

    def fit(self, arr: Sequence[Hashable]) -> "OneHotEncoder":
        self.columns = sorted(set(arr))
        self._index = {v: j for j, v in enumerate(self.columns)}
        return self

    def transform(self, arr: Sequence[Hashable]) -> list[list[int]]:
        k = len(self.columns)
        out = []
        for x in arr:
            row = [0] * k
            j = self._index.get(x)
            if j is None:
                if self.strict:
                    raise KeyError(f"unseen value: {x!r}")
            else:
                row[j] = 1
            out.append(row)
        return out


if __name__ == "__main__":
    assert one_hot([4, 1, 4, 7]) == [[0, 1, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]
    assert one_hot([]) == []
    assert one_hot([5, 5]) == [[1], [1]]
    assert one_hot([-3, 0, -3]) == [[1, 0], [0, 1], [1, 0]]
    cols, codes = encode_sparse([4, 1, 4, 7])
    assert cols == [1, 4, 7] and codes == [1, 0, 1, 2]
    assert decode([1, 4, 7], one_hot([4, 1, 4, 7])) == [4, 1, 4, 7]
    enc = OneHotEncoder().fit([2, 9])
    assert enc.transform([9, 3]) == [[0, 1], [0, 0]]
    print("ok")

Complexity

Let n = len(arr) and k = number of distinct values (k <= n).

  • one_hot: O(n + k log k) time to build the columns (set plus sort), then O(n·k) to write the matrix. Total O(n·k), which is the size of the output. Space O(n·k) for the output and O(k) for the index.
  • encode_sparse: O(n + k log k) time and O(n + k) space.
  • decode: O(n·k) time, because it reads every cell.
  • OneHotEncoder.fit: O(n + k log k). transform on m rows: O(m·k).

Test cases

  • [4, 1, 4, 7] -> columns [1, 4, 7] -> [[0,1,0],[1,0,0],[0,1,0],[0,0,1]]
  • [] -> [] (no rows and no columns)
  • [5] -> [[1]]
  • [5, 5, 5] -> [[1],[1],[1]] (one distinct value means one column)
  • [-3, 0, -3] -> columns [-3, 0] -> [[1,0],[0,1],[1,0]] (negatives sort first)
  • [3, 2, 1] -> columns [1, 2, 3] -> [[0,0,1],[0,1,0],[1,0,0]] (columns use sorted order, not order of appearance)
  • [10**9, -10**9] -> [[0,1],[1,0]] (the size of the values doesn't matter, only how many distinct values there are)
  • encode_sparse([4, 1, 4, 7]) -> ([1, 4, 7], [1, 0, 1, 2])
  • OneHotEncoder().fit([2, 9]).transform([9, 3]) -> [[0,1],[0,0]] (unseen 3 gives an all-zero row)

Follow-ups

  1. Decode back. Given the matrix and the column values, recover arr. Find the single 1 in each row and validate that the row is one-hot (exactly one 1, the rest 0s). That is decode above.
  2. Many distinct values / memory. With a million IDs the dense matrix doesn't fit in memory. Return (columns, codes) or a CSR-style triple (indptr, indices, data). Callers who need the dense row can build it on demand, for example with a generator that yields one row at a time.
  3. Fit vs transform (train/serve). The columns must come from the training data and stay fixed at serving time, or the features no longer line up with the model. Add a fitted encoder. Unseen values become an all-zero row, a reserved "unknown" column, or an error. Discuss which one the downstream model can handle.
  4. Other orderings or types. Strings or mixed types: sort with a key, or order columns by first appearance (scan once and assign the next column the first time a value is seen). Ordering by frequency (most common first) plus top-K with an "other" bucket limits the width.
  5. Multi-hot / NumPy. Each element can be a list of labels (genres per title). Set several 1s per row. In NumPy: cols, codes = np.unique(arr, return_inverse=True), then np.eye(len(cols), dtype=int)[codes] or m[np.arange(n), codes] = 1.

Interviewer expectations

Mid-level (L4)

  • Gets the column order right (sorted distinct values, not order of appearance) and uses a dict for the lookup, not list.index.
  • Handles the empty input and the case with only one distinct value without special cases breaking.
  • States O(n·k) and explains that it is bounded by the output size.

Senior (L5)

  • Writes decode with validation of malformed rows and points out the dense-vs-sparse trade-off without being prompted.
  • Separates fitting from transforming and has a clear policy for unseen values.
  • Mentions the vectorized NumPy form and when it pays off.

Staff (L6+)

  • Frames the problem as feature encoding for a model: training/serving consistency, versioning the column vocabulary, and the risk of high-cardinality explosion (top-K plus "other", or hashing).
  • Discusses how the encoding is stored and shipped (sparse formats, streaming rows) for very large inputs.
  • Keeps the core function small and testable, and pushes policy (ordering, unknowns) into configuration.

Sign in to comment, vote or remix.