Complexity Questions in ML Coding

Why Big-O Complexity Matters in ML Interviews

In machine learning coding interviews, writing working code is only half the requirement.

You must state the Time Complexity (Operations) and Space Complexity (RAM / VRAM Memory) of your algorithms in terms of key problem dimensions:

• $N$: Number of samples.
• $D$: Number of features or embedding dimensions.
• $T$: Sequence length.
• $K$: Number of clusters or neighbors.

┌─────────────────────────────────────────────────────────────┐
│ TIME COMPLEXITY:   Total scalar operations required.        │
│ SPACE COMPLEXITY:  Peak memory bytes allocated in RAM/VRAM. │
└─────────────────────────────────────────────────────────────┘

Essential ML Complexity Cheatsheet

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ ALGORITHM                │ TIME COMPLEXITY          │ SPACE COMPLEXITY         │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Matrix Multiply (N,M)x(M,P)│ O(N * M * P)           │ O(N * P)                 │
│ Linear Regression Fit    │ O(N * D + D^3)           │ O(N * D + D^2)           │
│ k-Means Iteration (K,N,D)│ O(N * K * D) per iter    │ O(N * D + K * D)         │
│ Decision Tree Training   │ O(N * D * log N)         │ O(Nodes + Depth)         │
│ Brute-Force k-NN Query   │ O(Q * N * D)             │ O(Q * N)                 │
│ Self-Attention (T, d)    │ O(T^2 * d)               │ O(T^2 + T * d)           │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

Worked Examples Explained

1. Matrix Multiplication: $A_{N \times M} \times B_{M \times P}$

To compute each entry in the output matrix of shape $(N, P)$, we perform $M$ multiplications. There are $N \times P$ entries, so total time is $\mathcal{O}(N \cdot M \cdot P)$.

2. Normal Equation Linear Regression: $w = (X^T X)^{-1} X^T y$

• $X^T X$ multiplication: $\mathcal{O}(D^2 N)$ time.
• Inverting $D \times D$ matrix: $\mathcal{O}(D^3)$ time.
• Total time: $\mathcal{O}(D^2 N + D^3)$. Explodes when feature dimension $D > 10,000$!

3. Self-Attention Mechanism: $\text{Softmax}(Q K^T / \sqrt{d_k}) V$

• $Q K^T$ where $Q, K \in \mathbb{R}^{T \times d}$: Matrix multiplication produces a $T \times T$ similarity matrix taking $\mathcal{O}(T^2 \cdot d)$ time.
• Storing the $T \times T$ attention weight matrix takes $\mathcal{O}(T^2)$ memory, creating the quadratic sequence length bottleneck!

Say this out loud

Big-O complexity analysis quantifies scalability for machine learning algorithms. Matrix multiplication takes O(N M P) time. Normal equation linear regression scales cubic O(D^3) with feature dimensions. Self attention requires O(T^2 d) time and O(T^2) memory, creating quadratic sequence length bottlenecks.

Followups to expect

1. What is the time complexity of sorting vs argpartition in NumPy? Full sorting (np.argsort) takes $\mathcal{O}(N \log N)$ time. Partial sorting (np.argpartition) extracts top $K$ smallest elements in linear $\mathcal{O}(N)$ average time.
2. How does Strassen's Algorithm improve matrix multiplication complexity? Strassen's algorithm reduces matrix multiplication complexity from $\mathcal{O}(N^3)$ down to approximately $\mathcal{O}(N^{2.81})$ using sub-matrix block decompositions.