Coding questions · Amazon
Design 2D DP on a DAG
Design 2D DP on a DAG evaluates algorithm design, data structures, correctness, complexity, edge cases, and implementation details in a realistic interview setting. A strong answer states assumptions, handles edge cases, explains trade-offs, and shows how to validate the result clearly.
Given a directed acyclic graph G(V, E) with N nodes, M edges, a designated source s, and an integer L, design a 2D dynamic program dp[u][k] that computes, for all nodes u and 0 ≤ k ≤ L, the number of distinct paths from s to u using exactly k edges. Specify the state definition, base cases, transition function, and iteration order using a topological sort. Analyze time and space complexity, describe how to reduce space via rolling arrays, and discuss adaptations for weighted edges (e.g., shortest/longest path with at most/exactly k edges) and for counting paths to a target t.
What a Strong Answer Covers
- A clear algorithm with the right data structures and enough pseudocode or code-level detail to implement it.
- A correctness argument that explains why the algorithm covers all required cases.
- Time and space complexity, plus at least one alternative approach when relevant.
- Focused tests for normal cases, edge cases, and failure modes.
Follow-up Questions
- How would the approach change if the input were streaming or too large for memory?
- What invariants would you assert in production code?
- Which tests would catch off-by-one, duplicate, or tie-breaking bugs?
Count DAG Paths with Exactly K Edges (2D DP)
Given a directed acyclic graph (DAG) with n nodes labeled 0..n-1, a list of directed edges where each edges[i] = [u, v] represents an edge u -> v, a designated source node, and an integer L, compute a 2D DP table dp where dp[u][k] is the number of distinct paths from source to node u that use exactly k edges, for every node u and every k with 0 <= k <= L.
Return dp as an n x (L+1) 2D list (row u, column k).
Use a topological order so each node is finalized before its successors are relaxed. Base case: dp[source][0] = 1, all other entries 0. Transition: for each node u in topological order and each edge u -> v, add dp[u][k] into dp[v][k+1] for 0 <= k < L.
Example: n=4, edges=[[0,1],[0,2],[1,3],[2,3]], source=0, L=2. There are two distinct 2-edge paths from 0 to 3 (0->1->3 and 0->2->3), so dp[3][2] = 2.
Assumptions: the graph is guaranteed acyclic; node indices are valid; counts fit in a 64-bit integer for the given test sizes.
Constraints
- 1 <= n <= 10^4
- 0 <= L <= 200
- 0 <= edges.length <= 5 * 10^4
- Each edge is [u, v] with 0 <= u, v < n; the graph is acyclic.
- 0 <= source < n
- Path counts fit in a signed 64-bit integer for the given inputs.
Hints
- Process nodes in topological order so dp[u][*] is fully computed before you push its contributions to successors.
- Base case is dp[source][0] = 1 (the zero-edge path from source to itself); every other entry starts at 0.
- Transition: for edge u -> v, every length-k path to u extends to a length-(k+1) path to v, so dp[v][k+1] += dp[u][k] for 0 <= k < L.
- To save memory when you only need exactly-L paths to one target, keep just two layers (k and k+1) and roll them — O(N) instead of O(N*L).
Solution
Solution Alignment
The prompt asks for an implementation-level answer. The safest way to present it is to define the state, maintain clear invariants, then walk through complexity and tests.
Problem Restatement
Given a directed acyclic graph G(V, E) with N nodes, M edges, a designated source s, and an integer L, design a 2D dynamic program dp[u][k] that computes, for all nodes u and 0 ≤ k ≤ L, the number of distinct paths from s to u using exactly k edges. Specify the state definition, base cases, transition function, and iteration order using a topological sort. Analyze time and space complexity, describe how to reduce space via rolling arrays, and discuss adaptations for weighted edges (e.g., shortest/longest path with at most/exactly k edges) and for counting paths to a target t.
Recommended Approach
Topologically sort the DAG. Define dp[u][k] as the number of paths from source s to node u using exactly k edges. Initialize dp[s][0] = 1 and all other states to 0. Iterate k from 0 to L-1; for every node u in topological order and edge u -> v, add dp[u][k] to dp[v][k+1].
Correctness
The implementation should maintain an invariant after each loop or operation that directly matches the problem statement. At termination, that invariant implies the returned value has considered every valid candidate exactly once, or has preserved the required data-structure state after every API call.
Complexity
Time: O(L * (N + M)) if each layer scans all edges. Space: O(N * (L + 1)); can be reduced to O(N) for a single target layer when earlier layers are no longer needed.
Edge Cases and Tests
L = 0, unreachable nodes, multiple paths to the same node, disconnected components, and very large counts requiring modulo arithmetic or big integers.