← Back to the 16 Patterns

ML and Data Processing

The other half of a data or ML interview: NumPy and pandas mental models, the statistics you will be asked to code from scratch, numerical stability, sampling, stream sketches, and the classic algorithms.

An ML or data-infrastructure loop rarely stops at the sixteen patterns. It adds two things: can you think in arrays rather than loops, and can you implement the maths without a library. This page covers both. Every pure-Python function here runs and is tested.

Where the work happens

The first question in any data problem is which layer you are working at. Choosing wrong costs an order of magnitude, and interviewers watch for it.

NumPy: thinking in arrays

The four things that define an array

import numpy as np

a = np.array([[1, 2, 3], [4, 5, 6]])       # O(n) to build
a.shape                                     # (2, 3)
a.dtype                                     # dtype('int64')

np.zeros((2, 3))                            # O(n)
np.arange(6).reshape(2, 3)                  # O(1) reshape: a VIEW, no copy
np.linspace(0.0, 1.0, 5)                    # 5 evenly spaced values

Vectorisation: the whole point

out = np.empty(len(a))
for i in range(len(a)):      # O(n) in PYTHON
    out[i] = a[i] * 2 + 1    # ~100x slower

out = a * 2 + 1              # O(n) in C,
                             # one pass, no
                             # Python objects

Broadcasting

import numpy as np

rows = np.zeros((3, 4))          # shape (3, 4)
per_column = np.array([1, 2, 3, 4])          # shape (4,)
rows + per_column                # OK: (3,4) and (4,) -> the row is reused

per_row = np.array([1, 2, 3])                # shape (3,)
# rows + per_row                 # ERROR: (3,4) vs (3,) do not line up
rows + per_row.reshape(3, 1)     # OK: (3,4) and (3,1) -> the column is reused

# The classic use: centre every column of a feature matrix.
features = np.random.rand(100, 5)            # O(n)
centred = features - features.mean(axis=0)   # mean is shape (5,), broadcast down

Views versus copies

import numpy as np

base = np.arange(10)

view = base[2:5]                 # O(1)  a VIEW. No data copied.
view[0] = 99                     # base[2] is now 99 as well

copied = base[2:5].copy()        # O(k)  independent
fancy = base[[2, 3, 4]]          # O(k)  fancy indexing ALWAYS copies
mask = base[base > 5]            # O(n)  boolean masking also copies

The operations worth knowing cold

pandas: thinking in tables

import pandas as pd

df = pd.DataFrame({"user": ["a", "b", "a"], "spend": [10, 20, 5]})

df.shape                    # (3, 2)
df.dtypes                   # per-column types; object means Python strings
df["spend"].sum()           # O(n) in C
df.head()                   # first 5 rows
df.info()                   # dtypes and MEMORY, the first thing to check

Selecting: loc versus iloc

The group-by mental model

import pandas as pd

df = pd.DataFrame({"user": ["a", "b", "a"], "spend": [10, 20, 5]})

df.groupby("user")["spend"].sum()               # O(n) split + O(n) reduce
df.groupby("user").agg(total=("spend", "sum"), n=("spend", "size"))

# Add a per-group value back onto every original row, without collapsing.
df["user_total"] = df.groupby("user")["spend"].transform("sum")

Joins

import pandas as pd

left = pd.DataFrame({"id": [1, 2], "x": ["a", "b"]})
right = pd.DataFrame({"id": [2, 3], "y": [9, 8]})

pd.merge(left, right, on="id", how="inner")     # O(n + m) via a hash join
pd.merge(left, right, on="id", how="left")      # keeps every left row
pd.concat([left, left])                         # stacks rows, O(n)

What makes pandas slow

Streaming and chunking

When the data does not fit in memory, the answer is almost always a generator pipeline. This is Guide 5 applied to real work, and it is pure standard library.

from collections.abc import Iterable, Iterator


def chunks(items: Iterable, size: int) -> Iterator[list]:
    """Yield lists of at most `size` items. O(size) memory, O(n) total.

    >>> list(chunks(range(5), 2))
    [[0, 1], [2, 3], [4]]
    """
    if size < 1:
        raise ValueError("size must be at least 1")

    batch: list = []
    for item in items:
        batch.append(item)          # O(1)
        if len(batch) == size:
            yield batch
            batch = []              # a NEW list: never yield the same object twice

    if batch:                       # the final short batch
        yield batch


def running_unique(items: Iterable) -> Iterator:
    """Drop duplicates while streaming. O(distinct) memory.

    >>> list(running_unique([1, 2, 1, 3, 2]))
    [1, 2, 3]
    """
    seen: set = set()
    for item in items:
        if item not in seen:        # O(1) average
            seen.add(item)
            yield item

The external-sort shape

The same shape answers “sort 100 GB with 8 GB of RAM”, “merge k sorted files”, and the shuffle stage of a map-reduce job.

Statistics, by hand

Expect to write these without a library, and to be asked about the numerical detail.

import math


def mean(values: list[float]) -> float:
    """Arithmetic mean. O(n) time, O(1) space.

    >>> mean([1.0, 2.0, 6.0])
    3.0
    """
    if not values:
        raise ValueError("mean of an empty sequence")
    return sum(values) / len(values)


def median(values: list[float]) -> float:
    """Middle value; the mean of the middle two when the count is even.

    O(n log n) as written. O(n) expected with quickselect, see Pattern 9.

    >>> median([3.0, 1.0, 2.0])
    2.0
    >>> median([1.0, 2.0, 3.0, 4.0])
    2.5
    """
    if not values:
        raise ValueError("median of an empty sequence")

    ordered = sorted(values)                    # O(n log n)
    mid = len(ordered) // 2
    if len(ordered) % 2 == 1:
        return ordered[mid]
    return (ordered[mid - 1] + ordered[mid]) / 2


def variance(values: list[float], *, sample: bool = True) -> float:
    """Spread around the mean. O(n) time, two passes.

    Args:
        sample: Divide by n-1 (Bessel's correction) for a sample, n for a
            whole population. Getting this wrong is a classic follow-up.

    >>> variance([2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0], sample=False)
    4.0
    """
    n = len(values)
    if n < 2:
        raise ValueError("variance needs at least two values")

    mu = mean(values)                                   # pass 1
    total = sum((v - mu) ** 2 for v in values)          # pass 2
    return total / (n - 1 if sample else n)


def standard_deviation(values: list[float], *, sample: bool = True) -> float:
    """Square root of the variance, back in the units of the data.

    >>> standard_deviation([2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0], sample=False)
    2.0
    """
    return math.sqrt(variance(values, sample=sample))

One pass, and numerically stable: Welford

class RunningStats:
    """Mean and variance in ONE pass, numerically stable. Welford's method.

    O(1) time per value, O(1) memory, so it works on an unbounded stream.

    >>> stats = RunningStats()
    >>> for x in [2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]:
    ...     stats.add(x)
    >>> round(stats.mean, 10)
    5.0
    >>> round(stats.population_variance, 10)
    4.0
    """

    def __init__(self) -> None:
        self.count = 0
        self.mean = 0.0
        self._sum_squared_deltas = 0.0

    def add(self, value: float) -> None:
        """Fold one more observation in. O(1)."""
        self.count += 1
        delta = value - self.mean
        self.mean += delta / self.count
        # Uses the mean BEFORE and AFTER the update, which is what keeps it stable.
        self._sum_squared_deltas += delta * (value - self.mean)

    @property
    def population_variance(self) -> float:
        return self._sum_squared_deltas / self.count if self.count else 0.0


def min_max_scale(values: list[float]) -> list[float]:
    """Rescale to [0, 1]. O(n).

    >>> min_max_scale([10.0, 20.0, 30.0])
    [0.0, 0.5, 1.0]
    """
    low, high = min(values), max(values)                # O(n) each
    if high == low:
        return [0.0] * len(values)                      # avoid dividing by zero
    span = high - low
    return [(v - low) / span for v in values]


def z_score(values: list[float]) -> list[float]:
    """Rescale to mean 0, standard deviation 1. O(n).

    >>> [round(z, 4) for z in z_score([2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0])]
    [-1.5, -0.5, -0.5, -0.5, 0.0, 0.0, 1.0, 2.0]
    """
    mu = mean(values)
    sigma = standard_deviation(values, sample=False)
    if sigma == 0.0:
        return [0.0] * len(values)
    return [(v - mu) / sigma for v in values]

Numerical stability

import math


def log_sum_exp(values: list[float]) -> float:
    """log(sum(exp(v))) without overflowing. O(n).

    exp(1000) is infinity, so the direct formula fails. Factoring out the
    largest term makes every exponent at most 0, so every exp is at most 1.

        log(sum(exp(v)))  ==  peak + log(sum(exp(v - peak)))

    >>> round(log_sum_exp([1000.0, 1000.0]), 6)
    1000.693147
    >>> round(log_sum_exp([0.0, 0.0]), 6)
    0.693147
    """
    if not values:
        return -math.inf

    peak = max(values)                                  # O(n)
    if peak == -math.inf:
        return -math.inf

    return peak + math.log(sum(math.exp(v - peak) for v in values))   # O(n)


def softmax(scores: list[float]) -> list[float]:
    """Turn scores into a probability distribution. O(n).

    Subtracting the maximum first changes nothing mathematically, because the
    factor cancels top and bottom, but it prevents overflow.

    >>> [round(p, 4) for p in softmax([1.0, 2.0, 3.0])]
    [0.09, 0.2447, 0.6652]
    >>> round(sum(softmax([1000.0, 1001.0])), 10)
    1.0
    """
    peak = max(scores)
    exponentials = [math.exp(s - peak) for s in scores]
    total = sum(exponentials)
    return [e / total for e in exponentials]


def sigmoid(x: float) -> float:
    """1 / (1 + exp(-x)), branching to keep the exponent negative. O(1).

    >>> round(sigmoid(0.0), 6)
    0.5
    >>> round(sigmoid(-1000.0), 6)
    0.0
    """
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    # For very negative x, exp(-x) would overflow. Use the equivalent form.
    positive = math.exp(x)
    return positive / (1.0 + positive)

Distance and similarity

import math


def dot(a: list[float], b: list[float]) -> float:
    """Dot product. O(n).

    >>> dot([1.0, 2.0], [3.0, 4.0])
    11.0
    """
    if len(a) != len(b):
        raise ValueError("vectors must have the same length")
    return sum(x * y for x, y in zip(a, b))


def euclidean(a: list[float], b: list[float]) -> float:
    """Straight-line distance, the L2 norm of the difference. O(n).

    >>> euclidean([0.0, 0.0], [3.0, 4.0])
    5.0
    """
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))


def manhattan(a: list[float], b: list[float]) -> float:
    """L1 distance: sum of absolute differences. O(n).

    >>> manhattan([0.0, 0.0], [3.0, 4.0])
    7.0
    """
    return sum(abs(x - y) for x, y in zip(a, b))


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """Cosine of the angle between two vectors, in [-1, 1]. O(n).

    Ignores magnitude entirely, which is why it is the default for text
    embeddings: a long document should not be "far" from a short one.

    >>> round(cosine_similarity([1.0, 0.0], [0.0, 1.0]), 6)
    0.0
    >>> round(cosine_similarity([1.0, 1.0], [2.0, 2.0]), 6)
    1.0
    """
    norm_a = math.sqrt(dot(a, a))
    norm_b = math.sqrt(dot(b, b))
    if norm_a == 0.0 or norm_b == 0.0:
        raise ValueError("cosine similarity is undefined for a zero vector")
    return dot(a, b) / (norm_a * norm_b)


def jaccard(a: set, b: set) -> float:
    """Overlap of two sets: intersection over union. O(min(|a|, |b|)).

    >>> jaccard({1, 2, 3}, {2, 3, 4})
    0.5
    """
    if not a and not b:
        return 1.0
    return len(a & b) / len(a | b)

Sampling

import bisect
import itertools
import random


def reservoir_sample(stream, k: int) -> list:
    """Take k items uniformly at random from a stream of UNKNOWN length.

    O(n) time, O(k) memory, and one pass. Item i survives with probability
    k/i at every later step, which works out to k/n overall.

    >>> reservoir_sample(range(3), 5)
    [0, 1, 2]
    >>> len(reservoir_sample(range(1000), 4))
    4
    """
    if k < 0:
        raise ValueError("k must be non-negative")

    reservoir: list = []
    for index, item in enumerate(stream):
        if index < k:
            reservoir.append(item)              # fill it first
        else:
            position = random.randint(0, index)  # inclusive on both ends
            if position < k:
                reservoir[position] = item       # evict a uniform survivor

    return reservoir


class WeightedSampler:
    """Draw an index with probability proportional to its weight.

    Prefix sums plus binary search: O(n) to build, O(log n) per draw.

    >>> sampler = WeightedSampler([0.0, 1.0, 0.0])
    >>> {sampler.draw() for _ in range(20)}     # only index 1 has any weight
    {1}
    """

    def __init__(self, weights: list[float]) -> None:
        if not weights or any(w < 0 for w in weights):
            raise ValueError("weights must be non-empty and non-negative")

        self._cumulative = list(itertools.accumulate(weights))   # O(n)
        self._total = self._cumulative[-1]
        if self._total <= 0:
            raise ValueError("weights must sum to a positive number")

    def draw(self) -> int:
        """One weighted draw. O(log n)."""
        target = random.random() * self._total
        return bisect.bisect_right(self._cumulative, target)


def train_test_split(rows: list, *, test_fraction: float = 0.2, seed: int = 0):
    """Shuffle then cut. O(n).

    A fixed seed makes the split reproducible, which matters more than it
    sounds: without it, two runs are not comparable.

    >>> train, test = train_test_split(list(range(10)), test_fraction=0.3)
    >>> len(train), len(test)
    (7, 3)
    >>> sorted(train + test) == list(range(10))
    True
    """
    if not 0.0 <= test_fraction < 1.0:
        raise ValueError("test_fraction must be in [0, 1)")

    shuffled = list(rows)
    random.Random(seed).shuffle(shuffled)        # O(n), Fisher-Yates
    cut = int(len(shuffled) * (1.0 - test_fraction))
    return shuffled[:cut], shuffled[cut:]

Sketches for streams

import hashlib


class BloomFilter:
    """Probabilistic set membership in a fixed bitmap.

    No false negatives: if it says "absent", the item was never added.
    False positives are possible, and their rate rises as you add items.

    Uses hashlib rather than the built-in hash(), because hash() on strings
    is randomised per process, so results would differ between runs.

    >>> bloom = BloomFilter()
    >>> bloom.add("apple")
    >>> "apple" in bloom
    True
    >>> "pear" in bloom
    False
    """

    def __init__(self, bits: int = 4096, hashes: int = 4) -> None:
        self._bits = bits
        self._hashes = hashes
        self._bitmap = bytearray((bits + 7) // 8)      # O(bits) memory, fixed

    def _positions(self, item: str):
        """Derive several independent positions from one digest. O(hashes)."""
        digest = hashlib.sha256(item.encode()).digest()
        for i in range(self._hashes):
            chunk = int.from_bytes(digest[i * 4 : i * 4 + 4], "big")
            yield chunk % self._bits

    def add(self, item: str) -> None:
        """O(hashes), independent of how many items are already in it."""
        for position in self._positions(item):
            self._bitmap[position // 8] |= 1 << (position % 8)

    def __contains__(self, item: str) -> bool:
        """O(hashes). True means "probably present", False means "definitely not"."""
        return all(
            self._bitmap[position // 8] & (1 << (position % 8))
            for position in self._positions(item)
        )


class CountMinSketch:
    """Approximate frequency counts in fixed memory.

    Every cell is shared by many keys, so collisions only ever ADD. Taking
    the minimum across rows discards the most-collided estimates, which is
    why the answer never underestimates.

    >>> sketch = CountMinSketch()
    >>> for _ in range(5):
    ...     sketch.add("clicks")
    >>> sketch.estimate("clicks")
    5
    >>> sketch.estimate("never-seen")
    0
    """

    def __init__(self, width: int = 512, depth: int = 4) -> None:
        self._width = width
        self._depth = depth
        self._table = [[0] * width for _ in range(depth)]     # O(width * depth)

    def _positions(self, key: str):
        for row in range(self._depth):
            digest = hashlib.sha256(f"{row}:{key}".encode()).digest()
            yield row, int.from_bytes(digest[:4], "big") % self._width

    def add(self, key: str, count: int = 1) -> None:
        """O(depth), regardless of stream length."""
        for row, column in self._positions(key):
            self._table[row][column] += count

    def estimate(self, key: str) -> int:
        """O(depth). Never below the true count, sometimes above."""
        return min(self._table[row][column] for row, column in self._positions(key))

Classic algorithms to implement

import random
from collections import Counter


def knn_classify(
    train: list[list[float]], labels: list, query: list[float], k: int
):
    """k-nearest-neighbours vote. O(n * d) to score, O(n log n) to rank.

    There is no training step: all the work happens at query time, which is
    the trade-off to name when asked.

    >>> points = [[0.0], [1.0], [10.0], [11.0]]
    >>> knn_classify(points, ["low", "low", "high", "high"], [9.5], k=3)
    'high'
    """
    if not 1 <= k <= len(train):
        raise ValueError(f"k must be in 1..{len(train)}")

    # For large n, a size-k heap is O(n log k) and beats the full sort.
    ranked = sorted(range(len(train)), key=lambda i: euclidean(train[i], query))
    votes = Counter(labels[i] for i in ranked[:k])
    return votes.most_common(1)[0][0]


def kmeans(
    points: list[list[float]], k: int, *, iterations: int = 20, seed: int = 0
) -> tuple[list[list[float]], list[int]]:
    """Lloyd's algorithm. O(iterations * n * k * d).

    Returns the centroids and the cluster index of each point. It converges
    to a LOCAL optimum, so the initial choice matters; real implementations
    use k-means++ seeding and several restarts.

    >>> data = [[0.0], [1.0], [10.0], [11.0]]
    >>> centroids, assignment = kmeans(data, 2)
    >>> sorted(sorted(c) for c in centroids)
    [[0.5], [10.5]]
    """
    if not 1 <= k <= len(points):
        raise ValueError(f"k must be in 1..{len(points)}")

    rng = random.Random(seed)
    centroids = [list(p) for p in rng.sample(points, k)]
    assignment = [0] * len(points)
    previous: list[int] | None = None

    for _ in range(iterations):
        # Assign: each point joins its nearest centroid. O(n * k * d)
        assignment = [
            min(range(k), key=lambda c: euclidean(point, centroids[c]))
            for point in points
        ]

        if assignment == previous:
            break                       # converged, nothing will change again
        previous = assignment

        # Update: each centroid moves to the mean of its members. O(n * d)
        for cluster in range(k):
            members = [p for p, a in zip(points, assignment) if a == cluster]
            if members:
                centroids[cluster] = [sum(axis) / len(members) for axis in zip(*members)]

    return centroids, assignment


def linear_regression_gd(
    xs: list[float],
    ys: list[float],
    *,
    learning_rate: float = 0.05,
    epochs: int = 2000,
) -> tuple[float, float]:
    """Fit y = w*x + b by gradient descent. O(epochs * n).

    The loss is mean squared error. Its gradients are
        d/dw = 2/n * sum((pred - y) * x)
        d/db = 2/n * sum(pred - y)

    >>> w, b = linear_regression_gd([0.0, 1.0, 2.0, 3.0], [1.0, 3.0, 5.0, 7.0])
    >>> round(w, 3), round(b, 3)
    (2.0, 1.0)
    """
    if len(xs) != len(ys) or not xs:
        raise ValueError("xs and ys must be non-empty and the same length")

    n = len(xs)
    w = b = 0.0

    for _ in range(epochs):
        errors = [w * x + b - y for x, y in zip(xs, ys)]        # O(n)
        gradient_w = 2.0 * sum(e * x for e, x in zip(errors, xs)) / n
        gradient_b = 2.0 * sum(errors) / n
        w -= learning_rate * gradient_w
        b -= learning_rate * gradient_b

    return w, b


def linear_regression_closed_form(
    xs: list[float], ys: list[float]
) -> tuple[float, float]:
    """The exact answer in one pass, no iteration. O(n).

    w = covariance(x, y) / variance(x), and b follows from the means.

    >>> w, b = linear_regression_closed_form([0.0, 1.0, 2.0, 3.0], [1.0, 3.0, 5.0, 7.0])
    >>> round(w, 10), round(b, 10)
    (2.0, 1.0)
    """
    n = len(xs)
    mean_x, mean_y = sum(xs) / n, sum(ys) / n

    covariance = sum((x - mean_x) * (y - mean_y) for x, y in zip(xs, ys))
    spread = sum((x - mean_x) ** 2 for x in xs)
    if spread == 0.0:
        raise ValueError("all x values are identical, the slope is undefined")

    w = covariance / spread
    return w, mean_y - w * mean_x

Evaluation metrics

def confusion(y_true: list[int], y_pred: list[int]) -> tuple[int, int, int, int]:
    """Return (true positives, false positives, false negatives, true negatives).

    Labels are 0 or 1. O(n).

    >>> confusion([1, 1, 0, 0], [1, 0, 1, 0])
    (1, 1, 1, 1)
    """
    if len(y_true) != len(y_pred):
        raise ValueError("label lists must be the same length")

    tp = fp = fn = tn = 0
    for actual, predicted in zip(y_true, y_pred):
        if predicted == 1 and actual == 1:
            tp += 1
        elif predicted == 1:
            fp += 1
        elif actual == 1:
            fn += 1
        else:
            tn += 1

    return tp, fp, fn, tn


def precision_recall_f1(
    y_true: list[int], y_pred: list[int]
) -> tuple[float, float, float]:
    """Precision, recall and their harmonic mean. O(n).

    Precision: of what I flagged, how much was right. Punishes false alarms.
    Recall:    of what was there, how much did I find. Punishes misses.

    >>> precision_recall_f1([1, 1, 0, 0], [1, 0, 1, 0])
    (0.5, 0.5, 0.5)
    """
    tp, fp, fn, _ = confusion(y_true, y_pred)

    precision = tp / (tp + fp) if tp + fp else 0.0
    recall = tp / (tp + fn) if tp + fn else 0.0
    if precision + recall == 0.0:
        return precision, recall, 0.0

    f1 = 2 * precision * recall / (precision + recall)
    return precision, recall, f1


def roc_auc(y_true: list[int], scores: list[float]) -> float:
    """Area under the ROC curve, via the rank formula. O(n log n).

    AUC equals the probability that a random positive outranks a random
    negative, so it can be read straight off the ranks with no thresholds.

    >>> roc_auc([0, 0, 1, 1], [0.1, 0.2, 0.8, 0.9])
    1.0
    >>> roc_auc([0, 1], [0.9, 0.1])
    0.0
    """
    positives = sum(y_true)
    negatives = len(y_true) - positives
    if positives == 0 or negatives == 0:
        raise ValueError("AUC needs at least one positive and one negative")

    order = sorted(range(len(scores)), key=lambda i: scores[i])   # O(n log n)
    rank_sum = 0.0
    index = 0

    while index < len(order):
        # Tied scores must SHARE the average rank, or AUC is wrong.
        end = index
        while end + 1 < len(order) and scores[order[end + 1]] == scores[order[index]]:
            end += 1

        average_rank = (index + end) / 2 + 1        # ranks are 1-based
        for position in range(index, end + 1):
            if y_true[order[position]] == 1:
                rank_sum += average_rank

        index = end + 1

    return (rank_sum - positives * (positives + 1) / 2) / (positives * negatives)

Which metric to argue for

Complexity of the common operations