Computation Graph Pruning to Keep the Optimal Path

Given a computation graph that “looks like a tree” (assume it is a DAG rooted at a root node), each node is a neural-network operator (e.g., conv, activation). Given an input (or input spec), you need to: Select an optimal execution path among multiple alternative branches (optimality defined by a cost function). Prune branches that are not on the chosen optimal path (remove them or mark them unreachable). Requirements Design data structures to represent nodes and edges. Implement: computing the optimal path (e.g., minimum total cost / maximum score), returning the set of kept nodes and/or pruned nodes. Notes You may model a cost function as node.cost and path cost as the sum of node costs (and optionally edge costs). Constraints: N ≤ 1e5. Example (abstract) Root 0 branches into 0→1→3 and 0→2→4. If cost(1)+cost(3) < cost(2)+cost(4), keep {0,1,3} and prune {2,4}. Example Input 5 4 0 1 1 3 0 2 2 4 5 0 1 1 1 10 0 Output kept: 0 1 3 pruned: 2 4