DAG utility with add, validate and query

Question: DAG utility with add, validate and query

Companies: Notion

Difficulty: Medium

Problem statement

Build a small directed-acyclic-graph helper: add edges, reject ones that would create a cycle, and answer reachability or ordering queries over the graph. Reported on 1point3acres' Notion bank; reconstructed in our own words.

Clarify the setup

A helper for directed acyclic graphs, the kind behind formula dependencies, task blockers or synced blocks in a document tool:

• add_node(n), add_edge(u, v): u must come before v. Reject an edge that would create a cycle and leave the graph unchanged.
• has_path(u, v), descendants(n), ancestors(n).
• topological_order(): every node after all of its predecessors; deterministic (smallest name first among ready nodes).

Ask: self-loops (a cycle; reject), duplicate edges (ignore), adding an edge to an unknown node (create it, or raise; create it here), and how often edges are added versus queried. That last answer decides between the simple version and part 2.

Approach

Part 1: check reachability on every insert. Adding u → v creates a cycle exactly when v can already reach u. Run a DFS from v along out-edges; if it finds u, reject. Otherwise insert. Cost O(V + E) per insert in the worst case. Queries:

• has_path, descendants: DFS forward. ancestors: DFS over reversed edges, which we keep as preds for that reason.
• topological_order: Kahn's algorithm with a min-heap of ready nodes for a deterministic order, O((V + E) log V).

Part 2: keep a topological order up to date (Pearce and Kelly). If inserts are frequent and the graph is large, a full DFS per insert is wasteful. Maintain ord[n], a position in a valid topological order:

• If ord[u] < ord[v], the new edge already agrees with the order: no cycle is possible. Insert in O(1). In practice most edges are like this.
• Otherwise only nodes with positions between ord[v] and ord[u] can be involved. DFS forward from v but only through nodes with ord <= ord[u] (if it reaches u, it is a cycle), and backward from u through nodes with ord >= ord[v]. Then reassign the positions of just those visited nodes: backward set first, forward set after, reusing the same pool of position numbers.

The work is proportional to the affected region, not the whole graph, and topological_order() becomes a sort by ord (or O(V) if positions are kept dense).

Code

import heapq
from typing import Dict, Hashable, List, Set

Node = Hashable


class CycleError(ValueError):
    pass


class DAG:
    def __init__(self) -> None:
        self.succ: Dict[Node, Set[Node]] = {}
        self.preds: Dict[Node, Set[Node]] = {}

    def add_node(self, n: Node) -> None:
        self.succ.setdefault(n, set())
        self.preds.setdefault(n, set())

    def add_edge(self, u: Node, v: Node) -> None:
        self.add_node(u)
        self.add_node(v)
        if v in self.succ[u]:
            return
        if u == v or self.has_path(v, u):
            raise CycleError(f"{u} -> {v} would create a cycle")
        self.succ[u].add(v)
        self.preds[v].add(u)

    def remove_edge(self, u: Node, v: Node) -> None:
        self.succ.get(u, set()).discard(v)
        self.preds.get(v, set()).discard(u)

    def _reach(self, start: Node, edges: Dict[Node, Set[Node]]) -> Set[Node]:
        seen: Set[Node] = set()
        stack = list(edges.get(start, ()))
        while stack:
            n = stack.pop()
            if n not in seen:
                seen.add(n)
                stack.extend(edges[n])
        return seen

    def has_path(self, u: Node, v: Node) -> bool:
        if u == v:
            return u in self.succ
        return v in self._reach(u, self.succ)

    def descendants(self, n: Node) -> Set[Node]:
        return self._reach(n, self.succ)

    def ancestors(self, n: Node) -> Set[Node]:
        return self._reach(n, self.preds)

    def topological_order(self) -> List[Node]:
        indegree = {n: len(p) for n, p in self.preds.items()}
        ready = [n for n, d in indegree.items() if d == 0]
        heapq.heapify(ready)
        order = []
        while ready:
            n = heapq.heappop(ready)
            order.append(n)
            for m in self.succ[n]:
                indegree[m] -= 1
                if indegree[m] == 0:
                    heapq.heappush(ready, m)
        return order


class IncrementalDAG(DAG):
    """Part 2: maintains a topological order so most inserts are O(1)."""

    def __init__(self) -> None:
        super().__init__()
        self.ord: Dict[Node, int] = {}
        self.checks = 0                           # nodes visited by cycle checks

    def add_node(self, n: Node) -> None:
        if n not in self.ord:
            self.ord[n] = len(self.ord)
        super().add_node(n)

    def add_edge(self, u: Node, v: Node) -> None:
        self.add_node(u)
        self.add_node(v)
        if v in self.succ[u]:
            return
        if u == v:
            raise CycleError(f"{u} -> {v} would create a cycle")
        lb, ub = self.ord[v], self.ord[u]
        if lb < ub:                                   # order must change around this edge
            forward = self._bounded(v, self.succ, lambda n: self.ord[n] <= ub)
            if u in forward:
                raise CycleError(f"{u} -> {v} would create a cycle")
            backward = self._bounded(u, self.preds, lambda n: self.ord[n] >= lb)
            nodes = sorted(backward, key=self.ord.get) + sorted(forward, key=self.ord.get)
            slots = sorted(self.ord[n] for n in nodes)
            for n, slot in zip(nodes, slots):
                self.ord[n] = slot
        self.succ[u].add(v)
        self.preds[v].add(u)

    def _bounded(self, start: Node, edges: Dict[Node, Set[Node]], keep) -> Set[Node]:
        seen = {start}
        stack = [start]
        while stack:
            n = stack.pop()
            self.checks += 1
            for m in edges[n]:
                if m not in seen and keep(m):
                    seen.add(m)
                    stack.append(m)
        return seen

    def order(self) -> List[Node]:
        return sorted(self.ord, key=self.ord.get)

Complexity

Space is O(V + E) for both edge maps, plus O(V) for ord.

Test cases

def test_rejects_cycles_and_leaves_graph_unchanged():
    g = DAG()
    g.add_edge("a", "b")
    g.add_edge("b", "c")
    for u, v in [("c", "a"), ("b", "b"), ("c", "b")]:
        try:
            g.add_edge(u, v)
            assert False
        except CycleError:
            pass
    assert g.succ == {"a": {"b"}, "b": {"c"}, "c": set()}


def test_queries_and_order():
    g = DAG()
    for u, v in [("shop", "cook"), ("cook", "eat"), ("invite", "eat"), ("clean", "invite")]:
        g.add_edge(u, v)
    g.add_node("relax")
    assert g.has_path("shop", "eat") and not g.has_path("eat", "shop")
    assert g.descendants("clean") == {"invite", "eat"}
    assert g.ancestors("eat") == {"shop", "cook", "invite", "clean"}
    assert g.topological_order() == ["clean", "invite", "relax", "shop", "cook", "eat"]


def test_order_respects_every_edge():
    import random
    rng = random.Random(1)
    g = DAG()
    for _ in range(300):
        u, v = rng.randrange(40), rng.randrange(40)
        try:
            g.add_edge(u, v)
        except CycleError:
            pass
    pos = {n: i for i, n in enumerate(g.topological_order())}
    assert len(pos) == len(g.succ)
    assert all(pos[u] < pos[v] for u in g.succ for v in g.succ[u])

def test_incremental_matches_simple_version():
    import random
    rng = random.Random(2)
    a, b = DAG(), IncrementalDAG()
    for _ in range(2000):
        u, v = rng.randrange(80), rng.randrange(80)
        ra = rb = "ok"
        try:
            a.add_edge(u, v)
        except CycleError:
            ra = "cycle"
        try:
            b.add_edge(u, v)
        except CycleError:
            rb = "cycle"
        assert ra == rb
    pos = b.ord
    assert all(pos[u] < pos[v] for u in b.succ for v in b.succ[u])
    assert sorted(pos.values()) == list(range(len(pos)))


def test_edges_in_order_are_constant_time():
    g = IncrementalDAG()
    for i in range(1000):
        g.add_node(i)
    for i in range(999):
        g.add_edge(i, i + 1)                      # agrees with insertion order
    assert g.checks == 0
    try:
        g.add_edge(999, 0)
        assert False
    except CycleError:
        pass

Follow-ups

1. Removing edges. Removal never creates a cycle and the current order stays valid, so it is O(1) in both versions.
2. Batch inserts. Adding many edges at once (importing a project) is cheaper as one pass: add them all, run Kahn's algorithm once, and if nodes remain unprocessed, report the cycle (the leftover nodes contain it) and roll the batch back.
3. Explaining a cycle. Users need to see why an edge is rejected: keep parent pointers in the DFS from v and return the path v → … → u, so the message reads "A depends on B, which depends on C, which depends on A".
4. Reachability queries at scale. For frequent has_path on a big, mostly static graph, precompute reachability with interval labels on a spanning tree plus exceptions, or cache transitive closures for hot nodes and invalidate on edge changes.
5. Persistence and concurrency. Store edges as rows with a unique constraint; serialize inserts per graph (one writer) so two concurrent inserts cannot each pass the cycle check and together form a cycle.

Interviewer expectations

Mid-level (L4)

• Correct cycle check on insert (reachability from v to u) and a working topological sort.
• Handles self-loops, duplicate edges and unknown nodes.
• Tests that rejected edges leave the graph unchanged.

Senior (L5)

• Keeps reverse edges for ancestor queries, deterministic ordering with a heap, and clear complexity per operation.
• Verifies the order against every edge in a randomized test.
• Explains the cycle to the user with the offending path.

Staff (L6+)

• Maintains a topological order incrementally so most inserts are O(1) and others touch only the affected region, and tests it against the simple version.
• Handles batch imports, concurrent writers and persistence.
• Picks between per-insert checks, incremental ordering and precomputed reachability from the workload.