Your Compiler

DSA for AI Engineer Interviews

Not a general DSA list. This is scoped to what actually gets asked in AI/ML engineering loops — including the implementation questions no general list covers, and an explicit list of what you can skip.

72 nodes across 5 sections.

0 of 21 high-frequency done
HighMediumLowSkip
0/11 high

The single most common tool in a screening round. Most 'medium' problems are a hashmap in disguise.

Frequency countingToken counts, vocab building, label distributions.
Dedup with a seen-setDeduplicating a corpus before embedding is a real task.
Grouping by computed keyBucketing samples by class, length, or shard.
Index mappingToken↔ID maps are exactly this.
Complement lookup (two-sum family)

Directly adjacent to retrieval. If you're interviewing for anything touching search or RAG, expect this.

Top-K elementsThis IS vector search's final step.
K closest pointsNearest-neighbour by another name.
Streaming top-K (bounded heap)Retrieval over a corpus too large to sort.
K-way mergeMerging results across index shards.
Two-heap running median

Cheap to learn, shows up constantly, and maps onto chunking work.

Fixed-size windowFixed-size chunking with stride.
Variable window (expand/shrink)Token-budget-aware chunking.
Opposite-end pointers
Fast/slow pointers

Rarely the whole question, almost always part of it.

Sorting by a computed keyRanking by score, then tie-breaking by recency.
Stability and tie-breakingReranking must not scramble equal scores.
Partial sort / quickselectYou need top-10, not a full sort of 10M.

Shows up more than you'd expect via chunk-overlap and span problems.

Merge overlapping intervalsMerging overlapping text spans after chunking.
Insert interval
Max concurrent intervalsConcurrency limits on inference workers.
0/1 high

Know both cold. Everything else in this section builds on them.

BFS by level
DFS, recursive and iterative
Grid as a graph
Cycle detectionValidating a DAG before you run it.

Underrated for this role. Any pipeline, agent graph, or dependency resolver is a topo sort.

Kahn's algorithm (in-degree)Scheduling stages in a training or inference DAG.
DFS-based topological order
Detecting an invalid DAGAgent tool graphs with circular dependencies.

Traversal and simple recursion. Balanced-tree internals are not asked here.

Pre/in/post-order traversal
Top-down vs bottom-up recursion
Depth and path problems

Occasionally asked, usually when the role touches tokenization or autocomplete.

Insert and search
Prefix matchingAutocomplete, and greedy longest-match tokenizers.
0/9 high

If you're interviewing anywhere near retrieval, assume at least one of these.

Cosine similarityThe most commonly asked AI-specific implementation question.
Dot product & Euclidean distance
L2 normalizationWhy normalized dot product == cosine is a standard follow-up.
Brute-force k-NN over N vectorsThen: 'now make it fast' → heap, then ANN.
Matrix multiply without numpy

Not deriving backprop — implementing the small, specific pieces correctly.

Numerically stable softmaxThe max-subtraction trick is the whole question. Overflow is the trap.
Scaled dot-product attentionWhy divide by sqrt(d_k) is the real question.
Layer normalization
Cross-entropy loss

Common at inference-focused companies and for anything agent-related.

Temperature scaling
Top-k sampling
Nucleus (top-p) samplingThe cumulative-probability cutoff trips people up.
Greedy vs beam search

String problems with an AI framing.

BPE merge loopA frequency-count + merge problem underneath.
Greedy longest-match tokenizer
Chunking with overlapSliding window, with a token budget instead of a fixed size.
Fitting context to a token budgetWhich chunks do you drop, and by what rule?

Asked more at ML-leaning companies than at pure engineering ones.

k-means from scratchInit, assign, update, convergence. Classic.
Precision, recall, F1 by hand
recall@k and MRRIf you claim retrieval experience, expect to implement these.
nDCG

The bridge between coding round and system design round. Often the last 10 minutes.

Token bucket rate limiterEvery LLM API client needs one. Very commonly asked.
LRU cacheThen: 'now cache embeddings with it'.
Dynamic batching with a max size and timeoutReal inference-server logic.
Retry with exponential backoff and jitterWhy jitter matters is the follow-up.

Know the shape: state, transition, base case, order.

1D DP
Grid DP
Memoization vs tabulation

Occasionally relevant — edit distance has a genuine NLP framing.

Edit distanceWord error rate, fuzzy matching against retrieved text.
Longest common subsequence

Real algorithms, real interviews somewhere — just not usually this one.

Segment trees / Fenwick trees
Bitmask DP
Digit DP
Kruskal / Prim (MST)
KMP / Rabin-Karp / Z-algorithm
Bellman-Ford / Floyd-Warshall
Implementing bubble/radix/bucket sort