Number of Ways to Assign Edge Weights I

medium tree math dfs parity

Problem

An undirected tree has n nodes labeled 1…n, rooted at node 1. Every edge starts at weight 0 and must be reassigned a weight of 1 or 2. Pick any node x at the maximum depth. Count the ways to weight the edges on the path from node 1 to x so its total cost is odd, modulo 10⁹ + 7.

Key observation: weight-2 edges never change parity, so only the count of weight-1 edges matters. If the path has d = maxDepth edges, the answer is the number of ways to pick an odd subset of those d edges to be 1 — which is 2^d−1.

Inputedges = [[1,2],[1,3],[3,4],[3,5]]

Output2

Max depth is 2 (nodes 4 and 5). On a 2-edge path the odd-cost assignments are (1,2) and (2,1): 2²⁻¹ = 2.

edges (u:v pairs, comma-separated)

def assignEdgeWeights(edges):
    MOD = 10**9 + 7
    n = len(edges) + 1
    g = [[] for _ in range(n + 1)]
    for u, v in edges:                  # build adjacency list
        g[u].append(v)
        g[v].append(u)
    depth = [0] * (n + 1)
    seen = [False] * (n + 1)
    seen[1] = True
    stack = [1]
    max_depth = 0
    while stack:                        # DFS from root = node 1
        node = stack.pop()
        for nb in g[node]:
            if not seen[nb]:
                seen[nb] = True
                depth[nb] = depth[node] + 1
                max_depth = max(max_depth, depth[nb])
                stack.append(nb)
    return pow(2, max_depth - 1, MOD)   # odd-subset count = 2^(d-1)

function assignEdgeWeights(edges) {
  const MOD = 1000000007n;
  const n = edges.length + 1;
  const g = Array.from({ length: n + 1 }, () => []);
  for (const [u, v] of edges) {         // build adjacency list
    g[u].push(v);
    g[v].push(u);
  }
  const depth = new Array(n + 1).fill(0);
  const seen = new Array(n + 1).fill(false);
  seen[1] = true;
  const stack = [1];
  let maxDepth = 0;
  while (stack.length) {                 // DFS from root = node 1
    const node = stack.pop();
    for (const nb of g[node]) {
      if (!seen[nb]) {
        seen[nb] = true;
        depth[nb] = depth[node] + 1;
        maxDepth = Math.max(maxDepth, depth[nb]);
        stack.push(nb);
      }
    }
  }
  let res = 1n;                          // 2^(maxDepth-1) mod MOD
  for (let i = 0; i < maxDepth - 1; i++) res = (res * 2n) % MOD;
  return Number(res);
}

int assignEdgeWeights(int[][] edges) {
    final long MOD = 1_000_000_007L;
    int n = edges.length + 1;
    List<List<Integer>> g = new ArrayList<>();
    for (int i = 0; i <= n; i++) g.add(new ArrayList<>());
    for (int[] e : edges) {              // build adjacency list
        g.get(e[0]).add(e[1]);
        g.get(e[1]).add(e[0]);
    }
    int[] depth = new int[n + 1];
    boolean[] seen = new boolean[n + 1];
    seen[1] = true;
    Deque<Integer> stack = new ArrayDeque<>();
    stack.push(1);
    int maxDepth = 0;
    while (!stack.isEmpty()) {           // DFS from root = node 1
        int node = stack.pop();
        for (int nb : g.get(node)) {
            if (!seen[nb]) {
                seen[nb] = true;
                depth[nb] = depth[node] + 1;
                maxDepth = Math.max(maxDepth, depth[nb]);
                stack.push(nb);
            }
        }
    }
    long res = 1;                        // 2^(maxDepth-1) mod MOD
    for (int i = 0; i < maxDepth - 1; i++) res = res * 2 % MOD;
    return (int) res;
}

int assignEdgeWeights(vector<vector<int>>& edges) {
    const long long MOD = 1000000007LL;
    int n = edges.size() + 1;
    vector<vector<int>> g(n + 1);
    for (auto& e : edges) {             // build adjacency list
        g[e[0]].push_back(e[1]);
        g[e[1]].push_back(e[0]);
    }
    vector<int> depth(n + 1, 0);
    vector<char> seen(n + 1, 0);
    seen[1] = 1;
    vector<int> stack = {1};
    int maxDepth = 0;
    while (!stack.empty()) {             // DFS from root = node 1
        int node = stack.back(); stack.pop_back();
        for (int nb : g[node]) {
            if (!seen[nb]) {
                seen[nb] = 1;
                depth[nb] = depth[node] + 1;
                maxDepth = max(maxDepth, depth[nb]);
                stack.push_back(nb);
            }
        }
    }
    long long res = 1;                   // 2^(maxDepth-1) mod MOD
    for (int i = 0; i < maxDepth - 1; i++) res = res * 2 % MOD;
    return (int) res;
}

Explanation

The problem mixes a tree traversal with a small counting argument. The traversal answers "how long is the path to the deepest node?", and the counting answers "how many weightings of that path are odd?".

Step 1 — find the maximum depth. Build an adjacency list and run a DFS (or BFS) from the root, node 1. The depth of the root is 0, and each child is one deeper than its parent. The largest depth reached is maxDepth; that is also the number of edges on the longest root-to-leaf path.

Step 2 — count odd-cost weightings. Each of the maxDepth edges is independently 1 or 2. A weight of 2 is even and never flips the total's parity, so the cost is odd exactly when an odd number of edges carry weight 1. The number of subsets of a d-element set with odd size is 2^d−1 — half of all 2^d subsets. So the answer is 2^maxDepth−1.

Why is it exactly half? Pair up subsets by toggling whether edge 1 is included; this pairing matches every even-sized subset with an odd-sized one, so the two classes are equal in size, each 2^d/2 = 2^d−1.

Because the exponent can be large (up to n−1 ≈ 10⁵), we reduce modulo 10⁹ + 7 using fast modular exponentiation (or a simple multiply-and-mod loop). All distinct deepest nodes share the same depth, so the choice of x never changes the answer.

Example: edges = [[1,2],[1,3],[3,4],[3,5]]. DFS gives depths 1, 1, 2, 2 for nodes 2, 3, 4, 5, so maxDepth = 2 and the answer is 2¹ = 2.

Time: O(n) Space: O(n)