Implement TF–IDF with sparse matrices

Implement TF–IDF from Scratch (Python + NumPy/SciPy)

You are given a list of documents (plain strings). Implement a TF–IDF vectorizer from scratch — no scikit-learn, no external NLP libraries — that turns the corpus into a sparse document–term matrix you could feed to a downstream model.

Your vectorizer must support the following behavior:

1. Tokenization and preprocessing. Lowercase the text, strip punctuation, and tokenize using a memory-efficient (streaming / generator-based) tokenizer. No external NLP libs.

2. Vocabulary and filtering. Support optional min_df and max_df document-frequency filtering. Each threshold may be an integer (an absolute number of documents) or a float in $(0, 1]$ (a fraction of documents). The vocabulary (the matrix columns) must be ordered lexicographically for deterministic output.

3. TF–IDF computation. Term frequency (TF) is the raw count of a term in a document. Use the smoothed IDF

   $\text{idf}_t = \log\!\left(\frac{1 + N}{1 + \text{df}_t}\right) + 1$

   where $N$ is the number of documents and $\text{df}_t$ is the number of documents containing term $t$. Output a SciPy CSR sparse matrix of shape (n_documents, vocab_size).

4. Transform behavior. At transform time, ignore out-of-vocabulary (OOV) tokens. Support optional L2 row normalization (each document vector scaled to unit length).

5. API surface. Expose a class with fit, transform, fit_transform, and inverse_transform(doc_index, k=None), where inverse_transform returns the top-$k$ terms for a document ranked by TF–IDF score (if k is None, return all non-zero terms sorted by score descending).

6. Complexity. Analyze the time and space complexity of fit and transform in terms of the number of documents, total tokens, and vocabulary size.

7. Unit test. Provide a unit test that checks correctness on a 3-document toy corpus containing repeated terms.

Hint: Where to start

Two passes over the corpus. Pass 1 (fit) computes document frequency: for each document, take the set of unique tokens (so a term repeated within one document still counts as df $+1$), increment a df dict, then build the lexicographically sorted vocabulary and the IDF vector. Pass 2 (transform) counts in-vocabulary tokens per document and assembles the sparse matrix.

Hint: Building the CSR matrix

Build the matrix in COO/CSR-input form directly: accumulate three parallel lists — data (the TF·IDF values), indices (column ids), and indptr (row boundaries) — then construct it in one shot via scipy.sparse.csr_matrix((data, indices, indptr), shape=...). Sort each row's column indices before appending so the CSR is in canonical order (important for deterministic, comparable output).

Hint: Threshold interpretation

Resolve min_df/max_df to absolute integer counts after you know $N$. A float threshold is a fraction of $N$ — ceil(min_df * N) for the lower bound and floor(max_df * N) for the upper bound is the conventional rounding. Then keep term $t$ iff $\text{min\_df}_{\text{abs}} \le \text{df}_t \le \text{max\_df}_{\text{abs}}$.

Hint: Inverse transform efficiency

For top-$k$ on a single CSR row, you only have the row's non-zero values. A full argsort is fine, but np.argpartition(-vals, k-1)[:k] selects the top-$k$ in $O(n)$ and you then sort just those $k$ — useful when a row has many non-zeros and $k$ is small.

Hint: Edge cases to watch

Empty documents (or all-OOV documents) produce all-zero rows — guard L2 normalization against dividing by a zero norm. An empty corpus ($N=0$) and a vocabulary that filters down to zero columns should both be handled without crashing. Smoothing in the IDF formula is what prevents division by zero when $\text{df}_t = N$.

Constraints & Assumptions

• Allowed dependencies: Python standard library, NumPy, and SciPy only — no scikit-learn, NLTK, spaCy, or other NLP libraries.
• Input is an iterable of strings. Assume it fits the streaming model: transform may re-iterate the corpus, so a re-iterable container (e.g. a list) is acceptable; a one-shot generator passed to a separate fit then transform would be consumed — call out this assumption.
• Handle empty documents gracefully (all-zero rows).
• Output must be deterministic: identical input yields identical column order and values across runs.
• Token definition can be assumed to be runs of lowercase alphanumeric characters ([a-z0-9]+) unless the interviewer specifies otherwise — state your choice.

Clarifying Questions to Ask

• How should tokens be defined — alphanumeric runs only, or keep apostrophes/hyphens (e.g. "don't", "state-of-the-art")? Should digits be kept or dropped?
• What TF variant is expected — raw count (as stated), log-scaled $1 + \log(\text{tf})$, or sublinear? (The prompt specifies raw count; confirm.)
• Should transform accept a fresh corpus (unseen documents), or only re-encode the fitted corpus? This drives whether OOV handling is exercised.
• Is idf defined exactly as given (the scikit-learn smoothed form), or should I support the unsmoothed variant $\log(N/\text{df})$ as well?
• What is the expected scale (number of documents, vocabulary size)? This decides whether per-document count caching in fit_transform is acceptable or whether memory must stay bounded.

What a Strong Answer Covers

• Correct, working code that compiles and passes the toy test: a class with fit / transform / fit_transform / inverse_transform, a streaming tokenizer, and a SciPy CSR output of the right shape.
• Faithful IDF and TF semantics: smoothed IDF exactly as specified, raw-count TF, and df computed from unique terms per document (not raw counts).
• Correct vocabulary handling: lexicographic column order, min_df/max_df resolved correctly for both int and float, OOV tokens dropped at transform time, L2 normalization applied per row and safe on zero rows.
• Sparse-structure fluency: building CSR via (data, indices, indptr) rather than densifying; awareness of canonical column ordering and why it matters.
• Complexity analysis that separates tokenization cost ($O(T)$), the $O(V \log V)$ vocabulary sort, and the $O(\text{nnz})$ matrix assembly — in terms of documents $D$, tokens $T$, vocabulary $V$, and non-zeros.
• Robustness: empty documents, empty corpus, all-OOV documents, and the max_df < min_df misconfiguration all handled deliberately.
• A meaningful unit test: asserts vocabulary order, specific IDF values, matrix shape/nnz, OOV behavior, normalized row norms $\approx 1$, and inverse_transform ranking.

Follow-up Questions

• How would you scale this to a corpus that does not fit in memory (tens of millions of documents)? Discuss the hashing trick (feature hashing) as an alternative to an explicit vocabulary, and the time/space trade-offs of dropping inverse_transform.
• The current transform requires the corpus to be re-iterable. How would you restructure fit_transform to make a single streaming pass, and what does that cost in peak memory?
• TF–IDF treats terms independently and ignores order. What would change if you needed n-grams (bigrams/trigrams), and how does that affect vocabulary size and sparsity?
• How would you compute cosine similarity between two documents directly from the CSR matrix, and why does L2 normalization make that a simple dot product?
• The smoothed IDF gives a term appearing in every document an IDF of $\log\!\big(\tfrac{1+N}{1+N}\big) + 1 = 1$, not $0$. Why might you prefer that over an unsmoothed IDF that zeroes such terms out?

Solution

Approach Overview

The vectorizer mirrors the structure of sklearn's TfidfVectorizer but is built from primitives:

• Tokenizer — a regex-based generator that streams lowercase alphanumeric tokens and discards punctuation. Being a generator keeps peak memory at one token, not one full token list per document.
• fit — one pass over the corpus to compute document frequency (df) from the unique terms in each document, then build a lexicographically sorted vocabulary and the smoothed IDF vector.
• transform — one pass to count in-vocabulary tokens per document, assemble a CSR matrix from (data, indices, indptr), apply TF·IDF weighting, and optionally L2-normalize each row.
• inverse_transform — read a single CSR row and return its highest-scoring terms.

The two-stage fit/transform split is what lets the vocabulary be fixed before any vector is produced, and what makes OOV handling well-defined (a token is OOV iff it is absent from the fitted vocabulary).

Key formula (smoothed IDF):

$\text{idf}_j = \log\!\left(\frac{1 + N}{1 + \text{df}_j}\right) + 1$

with $N$ = number of documents and $\text{df}_j$ = number of documents containing term $j$. The $+1$ inside both numerator and denominator (smoothing) prevents division by zero, and the trailing $+1$ guarantees every term keeps a non-zero weight even when it appears in every document.

---

Reference Implementation (Python + NumPy/SciPy only)

import re
import math
import numpy as np
from scipy.sparse import csr_matrix

class TfidfVectorizerScratch:
    """
    A TF–IDF vectorizer built from scratch using only Python, NumPy, and SciPy.

    Parameters
    ----------
    min_df : int or float, default=1
        - int   : keep terms with document frequency >= min_df.
        - float in (0, 1] : keep terms with df >= ceil(min_df * n_docs).
    max_df : int or float, default=1.0
        - int   : keep terms with document frequency <= max_df.
        - float in (0, 1] : keep terms with df <= floor(max_df * n_docs).
    normalize : bool, default=True
        If True, L2-normalize each row of the TF–IDF matrix.
    """

    _token_re = re.compile(r"[a-z0-9]+")  # lowercase alphanumeric runs

    def __init__(self, min_df=1, max_df=1.0, normalize=True):
        self.min_df = min_df
        self.max_df = max_df
        self.normalize = normalize

        # Learned after fit
        self.n_docs_ = None
        self.vocabulary_ = None   # term -> column index
        self.terms_ = None        # column index -> term (lexicographic)
        self.idf_ = None          # np.ndarray aligned with terms_
        self.df_ = None           # term -> df

        self._last_X = None        # cache for inverse_transform convenience

    # ------------------ Tokenization ------------------
    @classmethod
    def _iter_tokens(cls, text):
        """Streaming tokenizer: lowercase, strip punctuation, yield one token at a time."""
        for m in cls._token_re.finditer(text.lower()):
            yield m.group(0)

    # ------------------ Threshold resolution ------------------
    @staticmethod
    def _resolve_df_thresholds(min_df, max_df, n_docs):
        """Convert int/float thresholds into absolute integer df bounds."""
        if isinstance(min_df, float):
            if not (0.0 < min_df <= 1.0):
                raise ValueError("min_df float must be in (0, 1].")
            min_df_abs = math.ceil(min_df * n_docs)
        else:
            min_df_abs = int(min_df)

        if isinstance(max_df, float):
            if not (0.0 < max_df <= 1.0):
                raise ValueError("max_df float must be in (0, 1].")
            max_df_abs = math.floor(max_df * n_docs)
        else:
            max_df_abs = int(max_df)

        min_df_abs = max(min_df_abs, 1)
        if max_df_abs < min_df_abs:
            raise ValueError(
                f"max_df ({max_df_abs}) < min_df ({min_df_abs}); "
                "no term can satisfy both bounds."
            )
        return min_df_abs, max_df_abs

    # ------------------ Vocabulary / IDF (shared) ------------------
    def _build_vocab_and_idf(self, df, n_docs):
        """Given a df dict and n_docs, set terms_/vocabulary_/df_/idf_."""
        self.n_docs_ = n_docs
        if n_docs == 0:
            self.terms_ = []
            self.vocabulary_ = {}
            self.df_ = {}
            self.idf_ = np.zeros((0,), dtype=np.float64)
            return

        min_df_abs, max_df_abs = self._resolve_df_thresholds(
            self.min_df, self.max_df, n_docs
        )

        terms = [t for t, c in df.items() if min_df_abs <= c <= max_df_abs]
        terms.sort()  # lexicographic, deterministic column order

        self.terms_ = terms
        self.vocabulary_ = {t: i for i, t in enumerate(terms)}
        self.df_ = {t: df[t] for t in terms}

        idf = np.empty(len(terms), dtype=np.float64)
        for i, t in enumerate(terms):
            idf[i] = math.log((1 + n_docs) / (1 + df[t])) + 1.0
        self.idf_ = idf

    # ------------------ Fit ------------------
    def fit(self, corpus):
        """One pass to compute df from unique terms per document."""
        df = {}
        n_docs = 0
        for doc in corpus:
            n_docs += 1
            for tok in set(self._iter_tokens(doc)):   # unique terms => df, not tf
                df[tok] = df.get(tok, 0) + 1
        self._build_vocab_and_idf(df, n_docs)
        return self

    # ------------------ CSR assembly (shared) ------------------
    def _counts_to_csr(self, doc_count_iter, n_rows):
        """Build a CSR matrix from an iterable of {col_index: tf} dicts."""
        idf = self.idf_
        data, indices, indptr = [], [], [0]
        for per_col in doc_count_iter:
            for j in sorted(per_col):           # canonical column order per row
                data.append(per_col[j] * idf[j])
                indices.append(j)
            indptr.append(len(indices))

        X = csr_matrix(
            (
                np.asarray(data, dtype=np.float64),
                np.asarray(indices, dtype=np.int32),
                np.asarray(indptr, dtype=np.int32),
            ),
            shape=(n_rows, len(self.vocabulary_)),
            dtype=np.float64,
        )
        if self.normalize:
            self._l2_normalize_csr_inplace(X)
        self._last_X = X
        return X

    # ------------------ Transform ------------------
    def transform(self, corpus):
        """Encode documents with the fitted vocabulary/IDF. OOV tokens are dropped."""
        if self.vocabulary_ is None:
            raise RuntimeError("Call fit() before transform().")

        vocab = self.vocabulary_

        def doc_counts():
            for doc in corpus:
                per_col = {}
                for tok in self._iter_tokens(doc):
                    j = vocab.get(tok)
                    if j is not None:               # ignore OOV
                        per_col[j] = per_col.get(j, 0) + 1
                yield per_col

        # Materialize once so we know n_rows; counts are tiny vs. raw text.
        counts = list(doc_counts())
        return self._counts_to_csr(counts, n_rows=len(counts))

    # ------------------ fit_transform ------------------
    def fit_transform(self, corpus):
        """
        Single materialization of per-document counts to avoid re-iterating
        the corpus (useful when the corpus is a one-shot generator). Trades
        peak memory (O(sum of distinct terms per doc)) for one fewer pass.
        """
        df = {}
        doc_token_counts = []   # list of {token: tf}
        n_docs = 0
        for doc in corpus:
            n_docs += 1
            counts = {}
            for tok in self._iter_tokens(doc):
                counts[tok] = counts.get(tok, 0) + 1
            doc_token_counts.append(counts)
            for tok in counts:                       # unique terms => df
                df[tok] = df.get(tok, 0) + 1

        self._build_vocab_and_idf(df, n_docs)
        if n_docs == 0:
            X = csr_matrix((0, 0), dtype=np.float64)
            self._last_X = X
            return X

        vocab = self.vocabulary_

        def per_col_iter():
            for counts in doc_token_counts:
                per_col = {}
                for tok, tf in counts.items():
                    j = vocab.get(tok)
                    if j is not None:                # ignore terms filtered out by df
                        per_col[j] = tf
                yield per_col

        return self._counts_to_csr(per_col_iter(), n_rows=n_docs)

    # ------------------ Inverse transform ------------------
    def inverse_transform(self, doc_index, k=None, X=None, with_scores=True):
        """
        Top-k terms for a document, ranked by TF–IDF score (descending).
        If k is None, return all non-zero terms sorted by score.
        """
        if X is None:
            if self._last_X is None:
                raise RuntimeError("No matrix cached; pass X or call transform() first.")
            X = self._last_X
        if not isinstance(X, csr_matrix):
            raise TypeError("X must be a CSR matrix.")
        if not (0 <= doc_index < X.shape[0]):
            raise IndexError("doc_index out of range.")

        row = X.getrow(doc_index)
        if row.nnz == 0:
            return []

        vals, cols = row.data, row.indices
        if k is None or k >= len(vals):
            order = np.argsort(-vals)                 # full sort, descending
        else:
            top = np.argpartition(-vals, k - 1)[:k]   # O(n) top-k selection
            order = top[np.argsort(-vals[top])]       # then sort the k winners

        terms = self.terms_
        if with_scores:
            return [(terms[cols[i]], float(vals[i])) for i in order]
        return [terms[cols[i]] for i in order]

    # ------------------ Utilities ------------------
    @staticmethod
    def _l2_normalize_csr_inplace(X):
        """L2-normalize each row in place; zero rows are left untouched."""
        data, indptr = X.data, X.indptr
        for i in range(X.shape[0]):
            start, end = indptr[i], indptr[i + 1]
            if end > start:
                row = data[start:end]
                norm = math.sqrt(float(np.dot(row, row)))
                if norm > 0.0:
                    data[start:end] = row / norm

    def get_feature_names(self):
        return list(self.terms_) if self.terms_ is not None else []

Design notes

• df from unique terms. In fit/fit_transform we increment df from the set of a document's tokens, so a word repeated inside one document still counts as df $+1$. (TF, by contrast, uses raw counts.) Mixing these up is the most common correctness bug.
• Smoothing. log((1+N)/(1+df)) + 1 is finite and positive for all valid df, including df == N, so we never hit log(0) or a divide-by-zero.
• Canonical CSR. Within each row we sort column indices before appending. Sorted indices make the CSR canonical, which keeps equality/round-trip tests deterministic and lets downstream consumers assume sorted indices.
• OOV at two layers. A token can be missing because it was never seen (transform of a new corpus) or because df filtering removed it. Both are handled by the single vocab.get(tok) is None check.
• Zero-row safety. Empty or all-OOV documents become all-zero rows; L2 normalization skips them rather than dividing by zero.

---

Time and Space Complexity

Let

• $D$ = number of documents
• $T$ = total tokens scanned in one pass over the corpus
• $V$ = vocabulary size after df filtering
• $\text{nnz}$ = total non-zeros in the output matrix ($\sum_i r_i$, where $r_i$ = distinct in-vocab terms in document $i$)

fit

• Time: $O(T + V \log V)$ — $O(T)$ to tokenize and build per-document unique sets and the df dict; $O(V \log V)$ to sort the vocabulary lexicographically.
• Space: $O(V)$ for the df dict, vocabulary, and IDF vector (plus $O(\max_i r_i)$ for the transient per-document set).

transform

• Time: $O\!\big(T + \sum_i r_i \log r_i + \text{nnz}\big)$ — $O(T)$ to tokenize and count in-vocab tokens; $\sum_i r_i \log r_i$ to sort each row's columns into canonical order; $O(\text{nnz})$ to compute TF·IDF and assemble the CSR arrays (L2 normalization adds another $O(\text{nnz})$).
• Space: $O(\text{nnz})$ for the output matrix, plus $O(\max_i r_i)$ for the per-document count dict.

fit_transform caches per-document token counts to avoid a second pass, raising peak memory to $O(V + \text{nnz})$. For corpora that do not fit in memory, prefer fit(corpus) followed by transform(corpus) with a re-iterable corpus to keep the working set bounded.

---

Unit Test (3-document toy corpus)

Checks vocabulary order, smoothed IDF values, CSR shape/type, OOV handling, L2-normalized row norms, min_df filtering, and inverse_transform ranking. (math must be importable in the test module.)

import math
import unittest
import numpy as np
from scipy.sparse import csr_matrix

class TestTfidfVectorizerScratch(unittest.TestCase):
    def setUp(self):
        self.docs = [
            "The cat sat on the mat. The cat!",
            "Dog dog dog; cat sat.",
            "The dog and the cat played, and the dog slept.",
        ]

    def test_basic_fit_transform(self):
        vec = TfidfVectorizerScratch(min_df=1, max_df=1.0, normalize=True)
        X = vec.fit_transform(self.docs)

        expected_terms = sorted(
            {"and", "cat", "dog", "mat", "on", "played", "sat", "slept", "the"}
        )
        self.assertEqual(vec.get_feature_names(), expected_terms)
        self.assertEqual(X.shape, (3, len(expected_terms)))
        self.assertIsInstance(X, csr_matrix)

        # N = 3; df(cat) = 3 -> idf = log(4/4) + 1 = 1.0
        #        df(the) = 2 -> idf = log(4/3) + 1
        idf_by_term = dict(zip(vec.get_feature_names(), vec.idf_))
        self.assertAlmostEqual(idf_by_term["cat"], 1.0, places=7)
        self.assertAlmostEqual(idf_by_term["the"], math.log(4 / 3) + 1.0, places=7)

        # Every non-empty row is L2-normalized.
        for i in range(X.shape[0]):
            row = X.getrow(i)
            if row.nnz > 0:
                self.assertAlmostEqual(
                    float(np.sqrt((row.data ** 2).sum())), 1.0, places=7
                )

        # Doc 0: "the" (tf 3, idf log(4/3)+1) outranks "cat" (tf 2, idf 1.0).
        top2 = vec.inverse_transform(0, k=2, X=X, with_scores=False)
        self.assertEqual(top2, ["the", "cat"])

    def test_oov_ignored(self):
        vec = TfidfVectorizerScratch(normalize=False).fit(self.docs)
        X = vec.transform(["unseen tokens only"])
        self.assertEqual(X.shape[0], 1)
        self.assertEqual(X.getrow(0).nnz, 0)  # entirely OOV -> empty row

    def test_min_df_filtering(self):
        vec = TfidfVectorizerScratch(min_df=2, normalize=False)
        X = vec.fit_transform(self.docs)
        # df >= 2 keeps {cat, dog, sat, the}
        self.assertEqual(vec.get_feature_names(), sorted(["cat", "dog", "sat", "the"]))
        self.assertEqual(X.shape, (3, 4))

if __name__ == "__main__":
    unittest.main(argv=[""], exit=False)

These assertions all hold against the implementation above: the vocabulary is the nine lexicographically sorted terms, idf(cat) = 1.0 and idf(the) = log(4/3) + 1 ≈ 1.2877, every non-empty row has unit L2 norm, a fully-OOV document yields an empty row, and min_df=2 collapses the vocabulary to the four terms with df $\ge 2$.

---

Addressing the Follow-ups

• Out-of-memory corpora / hashing trick. Replace the explicit vocabulary with feature hashing: map each token to a column via hash(token) % n_features (a fixed, large dimension such as $2^{20}$). This needs no fit pass and no growing dict, so it streams in $O(T)$ time and $O(1)$ vocabulary memory — at the cost of hash collisions (two terms can share a column) and the loss of inverse_transform (the mapping is not invertible). IDF still needs document frequencies, so a single counting pass over hashed columns is required if you want the IDF weighting rather than raw hashed TF.
• True single-pass fit_transform. The version above caches per-document counts (peak memory $O(\text{nnz})$). A strictly streaming alternative computes df in one pass and TF·IDF in a second; you cannot do both in a single pass over a one-shot generator without buffering, because IDF depends on the global df that is only known after the whole corpus is seen.
• N-grams. Emit contiguous token windows (bigrams as "w_i w_{i+1}", etc.) from the tokenizer and treat each n-gram as a vocabulary term. Vocabulary size and nnz grow roughly linearly in the n-gram order, increasing both memory and the value of min_df filtering to prune rare combinations.
• Cosine similarity from CSR. With L2-normalized rows, $\cos(d_a, d_b) = \frac{x_a \cdot x_b}{\lVert x_a\rVert \,\lVert x_b\rVert} = x_a \cdot x_b$ because both norms are 1. So similarity is just X[a].dot(X[b].T) (a sparse dot product), and the full pairwise matrix is X @ X.T.
• Why smoothed IDF instead of zeroing common terms. Unsmoothed $\log(N/\text{df})$ assigns IDF $0$ to a term appearing in every document, deleting it from every vector. The smoothed form gives such a term IDF $1$, keeping a small but non-zero signal — robust when "every document" is an artifact of a small corpus, and it never produces $\log(0)$ for a term with $\text{df}=0$ at transform time.