38. Clone a Session-State Graph

Problem

Model a streaming session as a graph of state nodes, with edges representing “carries context into” transitions between related session states — for example a session that forks into a follow-up context, or references a prior one.

Before running an experiment against live session state, clone the whole reachable graph into a sandbox so mutations there never touch the original. The clone must not share any node objects with the original, and must preserve structure exactly, including cycles.

Source: src/38_clone_session_state_graph.py

class Node:
    def __init__(self, val: str, neighbors: list['Node'] | None = None): ...

def clone_session_graph(node: 'Node | None') -> 'Node | None':
    ...

Examples:
>>> root = build_graph({"ctx-a": ["ctx-b"], "ctx-b": ["ctx-a"]})
>>> clone = clone_session_graph(root)
>>> clone is not root
True
>>> clone.val, sorted(n.val for n in clone.neighbors)
('ctx-a', ['ctx-b'])

Step-by-Step Approach

1. Recognize this as the classic “clone graph” problem: traverse the reachable graph while building a parallel copy, being careful never to clone the same node twice even when it's reachable via multiple paths or a cycle.
2. Keep a map from original-node identity (id(node), since nodes aren't value-comparable) to its already-created clone.
3. Seed the map with a clone of the root, then run BFS (a queue) or DFS (recursion or a stack) over the original graph starting from the root.
4. For each original node popped, walk its neighbors: if a neighbor hasn't been cloned yet, clone it and enqueue it for its own traversal; either way, wire the current node's clone to that neighbor's clone.
5. Because the identity map is checked before cloning, a node reachable via a cycle back to itself is only ever cloned once — later visits just look it up.
6. Return the clone corresponding to the original root's identity; handle None input by returning None immediately.

The key insight is that the identity map is doing two jobs at once: it prevents infinite recursion on cycles, and it's exactly what stitches shared references back together correctly in the clone — a node visited from two different parents gets wired to the same single clone, not two divergent copies.

Reference solution

class Node:
    def __init__(self, val: str, neighbors: list['Node'] | None = None):
        self.val = val
        self.neighbors = neighbors if neighbors is not None else []


def clone_session_graph(node: 'Node | None') -> 'Node | None':
    # BFS with a visited map keyed by original node identity, O(V + E) time and space
    if node is None:
        return None
    clones: dict[int, Node] = {id(node): Node(node.val)}
    queue = [node]
    while queue:
        current = queue.pop(0)
        current_clone = clones[id(current)]
        for neighbor in current.neighbors:
            if id(neighbor) not in clones:
                clones[id(neighbor)] = Node(neighbor.val)
                queue.append(neighbor)
            current_clone.neighbors.append(clones[id(neighbor)])
    return clones[id(node)]


def build_graph(adjacency: dict[str, list[str]]) -> 'Node | None':
    if not adjacency:
        return None
    nodes = {val: Node(val) for val in adjacency}
    for val, neighbor_vals in adjacency.items():
        nodes[val].neighbors = [nodes[neighbor_val] for neighbor_val in neighbor_vals]
    return nodes[next(iter(adjacency))]


def graphs_structurally_equal(a: 'Node | None', b: 'Node | None') -> bool:
    if a is None or b is None:
        return a is None and b is None
    matched: dict[int, Node] = {}
    stack = [(a, b)]
    while stack:
        na, nb = stack.pop()
        if na.val != nb.val:
            return False
        if id(na) in matched:
            if matched[id(na)] is not nb:
                return False
            continue
        matched[id(na)] = nb
        if sorted(n.val for n in na.neighbors) != sorted(n.val for n in nb.neighbors):
            return False
        nb_by_val = {n.val: n for n in nb.neighbors}
        for na_neighbor in na.neighbors:
            stack.append((na_neighbor, nb_by_val[na_neighbor.val]))
    return True


def collect_node_ids(node: 'Node | None') -> set[int]:
    if node is None:
        return set()
    seen: set[int] = set()
    stack = [node]
    while stack:
        current = stack.pop()
        if id(current) in seen:
            continue
        seen.add(id(current))
        stack.extend(current.neighbors)
    return seen


TEST_CASES = [
    {
        "input": {"adjacency": {"solo": []}},
        "expected": {"adjacency": {"solo": []}},
    },
    {
        "input": {"adjacency": {"ctx-a": ["ctx-b"], "ctx-b": ["ctx-a"]}},
        "expected": {"adjacency": {"ctx-a": ["ctx-b"], "ctx-b": ["ctx-a"]}},
    },
    {
        "input": {
            "adjacency": {
                "s1": ["s2", "s4"],
                "s2": ["s1", "s3"],
                "s3": ["s2", "s4"],
                "s4": ["s1", "s3"],
            }
        },
        "expected": {
            "adjacency": {
                "s1": ["s2", "s4"],
                "s2": ["s1", "s3"],
                "s3": ["s2", "s4"],
                "s4": ["s1", "s3"],
            }
        },
    },
    {
        "input": {"adjacency": {}},
        "expected": {"adjacency": {}},
    },
    {
        "input": {"adjacency": {"x": ["y"], "y": ["x", "z"], "z": ["y"]}},
        "expected": {"adjacency": {"x": ["y"], "y": ["x", "z"], "z": ["y"]}},
    },
]


def main():
    for i, case in enumerate(TEST_CASES):
        adjacency = case["input"]["adjacency"]
        expected_adjacency = case["expected"]["adjacency"]
        print(f"Test {i}: clone_session_graph(build_graph({adjacency}))")
        original = build_graph(adjacency)

        if original is None:
            print("  original is None, cloning trivially")
            clone = clone_session_graph(original)
        else:
            clones: dict[int, Node] = {id(original): Node(original.val)}
            queue = [original]
            print(f"  BFS start: cloned root {original.val!r}")
            while queue:
                current = queue.pop(0)
                current_clone = clones[id(current)]
                for neighbor in current.neighbors:
                    if id(neighbor) not in clones:
                        clones[id(neighbor)] = Node(neighbor.val)
                        queue.append(neighbor)
                        print(f"    discovered {neighbor.val!r} via {current.val!r}, cloned")
                    current_clone.neighbors.append(clones[id(neighbor)])
                    print(f"    wired clone({current.val!r}) -> clone({neighbor.val!r})")
            clone = clones[id(original)]

        expected_root = build_graph(expected_adjacency)
        original_ids = collect_node_ids(original)
        clone_ids = collect_node_ids(clone)
        assert original_ids.isdisjoint(clone_ids), "clone shares Node identity with original"
        assert graphs_structurally_equal(clone, expected_root)
        print(f"PASSED: clone matches adjacency {expected_adjacency}")
    print(f"All {len(TEST_CASES)} test cases passed.")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• clones: dict[int, Node] — maps original node identity to its clone; the traversal's visited set and output builder in one
• id(node) — identity key, since Node objects don't define value equality
• queue.pop(0) — BFS frontier; a deque would be more efficient but a list works for small graphs
• current_clone.neighbors.append(clones[id(neighbor)]) — wires the clone graph's edges to mirror the original's

How to Recognize This Pattern