3558. Number of Ways to Assign Edge Weights I

LeetCode ↗

Problem Pattern Idea Example Code Complex Pitfalls

Problem Description

There is an undirected tree with n nodes labeled from 1 to n, rooted at node 1. The tree is represented by a 2D integer array edges of length n - 1, where edges[i] = [uᵢ, vᵢ] indicates that there is an edge between nodes uᵢ and vᵢ.

Initially, all edges have a weight of 0. You must assign each edge a weight of either 1 or 2.

The cost of a path between any two nodes u and v is the total weight of all edges in the path connecting them.

Select any one node x at the maximum depth. Return the number of ways to assign edge weights in the path from node 1 to x such that its total cost is odd.

Since the answer may be large, return it modulo 10⁹ + 7.

Note: Ignore all edges not in the path from node 1 to x.

In short, you are given a rooted tree and asked to focus on a path from the root (node 1) to a deepest node x. If this path contains d edges (where d is the maximum depth of the tree), each edge can independently be assigned a weight of either 1 or 2. The task is to count how many of the 2ᵈ possible weight assignments produce an odd total cost along that path, returning the result modulo 10⁹ + 7.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

How We Pick the Algorithm

Why DFS?

This problem maps to DFS through a short path in the full flowchart.

The problem provides a tree and asks for the maximum depth via DFS from the root, then applies a combinatorial formula (2^(d-1)) for the parity count.

Open in Flowchart

Show step-by-step reasoning

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

Yes: The problem provides a set of nodes connected by edges, which forms a graph structure.

Is it a tree?

Yes: The problem explicitly states there is an undirected tree with n nodes rooted at node 1, represented by n - 1 edges.

DFS

Yes: Since the structure is a tree, we use a Depth-First Search to traverse it. Here, we apply DFS to compute the maximum depth of the tree by recursively descending into each child and tracking the longest path from the root.

Conclusion: The flowchart guides us toward the Depth-First Search (DFS) pattern. We traverse the tree from the root to find the maximum depth d, then apply the combinatorial result that the number of odd-cost weight assignments along a path of d edges is 2^(d-1).

Intuition

The key observation is that the problem only cares about the path from the root (node 1) to a deepest node x. The actual identity of x does not matter — what matters is the length of that path, i.e., the number of edges on it. If the maximum depth of the tree is d, then the path we focus on contains exactly d edges.

Each of these d edges can be independently assigned a weight of either 1 or 2. We want to count the assignments whose total cost is odd.

Here comes the crucial simplification: the parity of the total cost depends only on how many edges get an odd weight. Since 2 is even and 1 is odd, an edge contributes to the parity only when it is assigned 1. So the total cost is odd exactly when we pick an odd number of edges to assign the weight 1. The edges assigned 2 never affect parity.

So the question reduces to a pure counting problem: out of d edges, in how many ways can we choose an odd-sized subset to receive weight 1 (the rest receive weight 2)?

For a set of d elements, the number of subsets with an odd size equals the number with an even size, and each is exactly half of the total 2^d subsets. This gives the well-known identity:

(number of odd-sized subsets) = 2^(d-1)

Therefore, once we use DFS to find the maximum depth d, the answer is simply 2^(d-1), computed with fast exponentiation under modulo 10^9 + 7.

Pattern Learn more about Tree, Depth-First Search and Math patterns.

Solution Approach

We use DFS + Mathematics to solve the problem in two stages.

Step 1: Build the adjacency list.

Since the tree is given as an edge list, we first construct an adjacency list g, where g[u] holds all the neighbors of node u. We allocate g with size n + 1 because nodes are labeled from 1 to n:

n = len(edges) + 1
g = [[] for _ in range(n + 1)]
for u, v in edges:
    g[u].append(v)
    g[v].append(u)

Step 2: Find the maximum depth using DFS.

We define a recursive function dfs(i, fa) that returns the depth of the subtree rooted at node i, where fa is the parent of i (used to avoid walking back up the tree). For each child j of i, we recurse and take the maximum of dfs(j, i) + 1. The + 1 accounts for the edge connecting i to its child:

def dfs(i: int, fa: int = 0) -> int:
    res = 0
    for j in g[i]:
        if j != fa:
            res = max(res, dfs(j, i) + 1)
    return res

d = dfs(1)

Calling dfs(1) returns d, the maximum depth measured in number of edges along the longest path from the root.

Step 3: Apply the combinatorial formula.

As reasoned in the intuition, the number of ways to assign weights so the total cost is odd equals the number of odd-sized subsets of the d edges, which is 2^(d-1). We compute this with fast exponentiation under the modulus 10^9 + 7:

return pow(2, d - 1, 10**9 + 7)

Python's built-in pow(base, exp, mod) performs modular exponentiation efficiently in O(log d) time.

Complexity Analysis:

Time complexity: O(n), where n is the number of nodes. The DFS visits every node exactly once, and the final modular exponentiation takes O(log d), which is dominated by the traversal.
Space complexity: O(n) for the adjacency list, plus O(n) for the recursion stack in the worst case (a skewed, chain-like tree).

Example Walkthrough

Let's trace through the solution with a concrete example.

Input:

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

This describes a tree with 5 nodes (n = 5). Let's visualize it:

Step 1: Build the adjacency list.

We process each edge and add both directions (since the tree is undirected):

Edge	Action
`[1, 2]`	`g[1] += [2]`, `g[2] += [1]`
`[1, 3]`	`g[1] += [3]`, `g[3] += [1]`
`[2, 4]`	`g[2] += [4]`, `g[4] += [2]`
`[4, 5]`	`g[4] += [5]`, `g[5] += [4]`

Resulting adjacency list:

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

Step 2: Find the maximum depth via DFS.

We call dfs(1, 0) and recurse, skipping the parent at each step:

dfs(1, 0):
  ├─ child 2: dfs(2, 1)
  │            ├─ child 4: dfs(4, 2)
  │            │            └─ child 5: dfs(5, 4) → returns 0  (leaf, no children except parent)
  │            │            so res = max(0, 0 + 1) = 1
  │            │            returns 1
  │            so res = max(0, 1 + 1) = 2
  │            returns 2
  │  res = max(0, 2 + 1) = 3
  ├─ child 3: dfs(3, 1) → returns 0  (leaf)
  │  res = max(3, 0 + 1) = 3
  returns 3

So the maximum depth is d = 3. The deepest node is x = 5, and the path 1 → 2 → 4 → 5 contains exactly 3 edges.

Step 3: Apply the combinatorial formula.

The path has 3 edges. Each edge can be assigned weight 1 (odd) or 2 (even). We want the total cost to be odd, which happens exactly when an odd number of edges receive weight 1.

Let's enumerate all 2³ = 8 assignments (edge weights along the path, where each entry is the weight given):

Edge weights	Total cost	Parity	Counts?
(1, 1, 1)	3	odd	✅
(1, 1, 2)	4	even	❌
(1, 2, 1)	4	even	❌
(2, 1, 1)	4	even	❌
(1, 2, 2)	5	odd	✅
(2, 1, 2)	5	odd	✅
(2, 2, 1)	5	odd	✅
(2, 2, 2)	6	even	❌

There are exactly 4 odd-cost assignments — these are precisely the ones that use an odd number of weight-1 edges (either one or three of them).

This matches the formula:

2^(d-1) = 2^(3-1) = 2^2 = 4

Output:

pow(2, 3 - 1, 10**9 + 7) = 4

The answer is 4, confirming that out of the 8 possible weight assignments along the deepest path, exactly half produce an odd total cost.

Solution Implementation

1from typing import List
2
3
4class Solution:
5    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
6        # Compute the maximum depth (longest path from the given node down the tree)
7        def dfs(node: int, parent: int = 0) -> int:
8            max_depth = 0
9            for child in graph[node]:
10                if child != parent:
11                    # Add 1 for the edge connecting node to its child
12                    max_depth = max(max_depth, dfs(child, node) + 1)
13            return max_depth
14
15        # Number of nodes is number of edges + 1
16        num_nodes = len(edges) + 1
17
18        # Build adjacency list (1-indexed, so size num_nodes + 1)
19        graph = [[] for _ in range(num_nodes + 1)]
20        for u, v in edges:
21            graph[u].append(v)
22            graph[v].append(u)
23
24        # Find the depth of the tree starting from root node 1
25        depth = dfs(1)
26
27        # The number of valid weight assignments making the root-to-deepest-leaf
28        # path sum odd is 2^(depth - 1), taken modulo 1e9 + 7
29        MOD = 10**9 + 7
30        return pow(2, depth - 1, MOD)
31```
32
33**Explanation of the approach:**
34
351. **Problem intuition:** Each edge can be assigned a weight of 1 or 2. We need the number of ways to assign weights so that the path sum from the root to the deepest node satisfies a parity condition (odd sum).
36
372. **Key insight:** For a path of length `d` (number of edges), the count of weight combinations producing an odd total is exactly `2^(d-1)`. This holds because exactly half of all `2^d` combinations yield an odd sum.
38
393. **DFS traversal:** The `dfs` function recursively explores the tree to find the longest root-to-leaf path (the tree's maximum depth in terms of edge count).
40
414. **Modular exponentiation:** `pow(2, depth - 1, MOD)` efficiently computes the result modulo `10^9 + 7`.
42
43**Alternative perspective (iterative BFS to avoid recursion depth limits):**
44
45For very deep trees, the recursive DFS could hit Python's recursion limit. A BFS-based level traversal would compute the maximum depth iteratively:
46
47```python3
48from collections import deque
49from typing import List
50
51
52class Solution:
53    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
54        num_nodes = len(edges) + 1
55        graph = [[] for _ in range(num_nodes + 1)]
56        for u, v in edges:
57            graph[u].append(v)
58            graph[v].append(u)
59
60        # BFS to find maximum depth (edge count) from root node 1
61        depth = 0
62        visited = [False] * (num_nodes + 1)
63        queue = deque([1])
64        visited[1] = True
65        while queue:
66            for _ in range(len(queue)):
67                node = queue.popleft()
68                for child in graph[node]:
69                    if not visited[child]:
70                        visited[child] = True
71                        queue.append(child)
72            depth += 1  # increment per level
73
74        # depth counts levels; edge count is depth - 1
75        MOD = 10**9 + 7
76        return pow(2, depth - 2, MOD)
77

1class Solution {
2    // Adjacency list representation of the tree
3    private List<Integer>[] graph;
4
5    /**
6     * Computes the number of ways to assign edge weights (each edge can be 1 or 2)
7     * such that the path from the root (node 1) to the deepest leaf has an odd total weight.
8     *
9     * The total weight parity along the longest path depends on the number of edges on it.
10     * Given 'd' edges on the longest path, exactly half of all 2^d assignments yield an odd sum,
11     * which equals 2^(d-1).
12     *
13     * @param edges the edge list of the tree
14     * @return the number of valid assignments modulo 1e9 + 7
15     */
16    public int assignEdgeWeights(int[][] edges) {
17        // Number of nodes equals number of edges plus one (tree property)
18        int nodeCount = edges.length + 1;
19
20        // Initialize adjacency list for nodes 1..nodeCount
21        graph = new List[nodeCount + 1];
22        Arrays.setAll(graph, index -> new ArrayList<>());
23
24        // Build the undirected tree
25        for (int[] edge : edges) {
26            int from = edge[0];
27            int to = edge[1];
28            graph[from].add(to);
29            graph[to].add(from);
30        }
31
32        // depth = number of edges on the longest root-to-leaf path
33        int maxDepth = dfs(1, 0);
34
35        // Result is 2^(maxDepth - 1) modulo 1e9 + 7
36        return (int) pow(2, maxDepth - 1, 1_000_000_007);
37    }
38
39    /**
40     * Depth-first search to find the maximum number of edges
41     * from the current node down to any leaf in its subtree.
42     *
43     * @param node   the current node being visited
44     * @param parent the parent of the current node (to avoid revisiting)
45     * @return the maximum edge-depth of the subtree rooted at 'node'
46     */
47    private int dfs(int node, int parent) {
48        int maxEdges = 0;
49        for (int child : graph[node]) {
50            // Skip the parent to traverse only downward in the tree
51            if (child != parent) {
52                maxEdges = Math.max(maxEdges, dfs(child, node) + 1);
53            }
54        }
55        return maxEdges;
56    }
57
58    /**
59     * Fast modular exponentiation: computes (base^exponent) % mod.
60     *
61     * @param base     the base value
62     * @param exponent the exponent value
63     * @param mod      the modulus
64     * @return (base^exponent) % mod
65     */
66    private long pow(long base, int exponent, int mod) {
67        long result = 1;
68        while (exponent > 0) {
69            // If the current lowest bit is set, multiply the result by the current base
70            if ((exponent & 1) != 0) {
71                result = result * base % mod;
72            }
73            // Square the base for the next bit
74            base = base * base % mod;
75            exponent >>= 1;
76        }
77        return result;
78    }
79}
80

1class Solution {
2public:
3    int assignEdgeWeights(vector<vector<int>>& edges) {
4        // Number of nodes equals number of edges plus one (tree property)
5        int nodeCount = edges.size() + 1;
6
7        // Build an adjacency list; index 0 is unused since nodes are 1-based
8        vector<vector<int>> adjacency(nodeCount + 1);
9        for (const auto& edge : edges) {
10            int from = edge[0];
11            int to = edge[1];
12            adjacency[from].push_back(to);
13            adjacency[to].push_back(from);
14        }
15
16        // Compute the maximum depth (longest path length from the root)
17        // using a recursive DFS that returns the height of the subtree.
18        auto dfs = [&](this auto&& dfs, int current, int parent) -> int {
19            int maxDepth = 0;
20            for (int next : adjacency[current]) {
21                if (next != parent) {
22                    // Add 1 for the edge connecting current to its child
23                    maxDepth = max(maxDepth, dfs(next, current) + 1);
24                }
25            }
26            return maxDepth;
27        };
28
29        // The number of valid weight assignments making the path's total
30        // weight odd is 2^(depth - 1), taken modulo 1e9 + 7.
31        const int MOD = 1000000007;
32        return power(2, dfs(1, 0) - 1, MOD);
33    }
34
35private:
36    // Fast exponentiation: computes (base^exponent) % mod iteratively.
37    long long power(long long base, int exponent, int mod) {
38        long long result = 1;
39        while (exponent > 0) {
40            // If the current lowest bit is set, multiply result by base
41            if (exponent & 1) {
42                result = result * base % mod;
43            }
44            // Square the base and shift exponent right by one bit
45            base = base * base % mod;
46            exponent >>= 1;
47        }
48        return result;
49    }
50};
51

1const MOD = 1_000_000_007;
2
3/**
4 * Fast modular exponentiation: computes (a^exponent) % modulo.
5 * Uses BigInt internally to avoid overflow during multiplication.
6 */
7function powMod(a: number, exponent: number, modulo: number): number {
8    let result = 1n;
9    let base = BigInt(a);
10    const m = BigInt(modulo);
11
12    while (exponent > 0) {
13        // If the current lowest bit is set, multiply it into the result
14        if (exponent & 1) {
15            result = (result * base) % m;
16        }
17        // Square the base for the next bit
18        base = (base * base) % m;
19        // Move to the next bit
20        exponent >>= 1;
21    }
22
23    return Number(result);
24}
25
26/**
27 * Assigns edge weights so that the sum of weights along the longest
28 * root-to-node path is odd, and returns the number of valid assignments.
29 *
30 * The answer equals 2^(maxDepth - 1) mod (10^9 + 7), where maxDepth is
31 * the number of edges on the longest path starting from node 1.
32 */
33function assignEdgeWeights(edges: number[][]): number {
34    // A tree with `n` nodes has exactly n - 1 edges
35    const n = edges.length + 1;
36
37    // Build adjacency list; nodes are 1-indexed, so size is n + 1
38    const graph: number[][] = Array.from({ length: n + 1 }, () => []);
39
40    for (const [u, v] of edges) {
41        graph[u].push(v);
42        graph[v].push(u);
43    }
44
45    /**
46     * Depth-first search returning the maximum number of edges
47     * from the current node down to the deepest descendant.
48     *
49     * @param node   the current node being visited
50     * @param parent the node we came from (to avoid revisiting)
51     */
52    const dfs = (node: number, parent: number): number => {
53        let maxDepth = 0;
54        for (const next of graph[node]) {
55            // Skip the parent to prevent cycling back
56            if (next !== parent) {
57                maxDepth = Math.max(maxDepth, dfs(next, node) + 1);
58            }
59        }
60        return maxDepth;
61    };
62
63    // Compute the longest path length (in edges) starting from root node 1,
64    // then count the valid weight assignments as 2^(length - 1).
65    const longestPath = dfs(1, 0);
66    return powMod(2, longestPath - 1, MOD);
67}
68

Time and Space Complexity

Time Complexity: O(n)

The algorithm performs a depth-first search (DFS) starting from node 1. During the traversal, each node is visited exactly once, and for each node, we iterate over its adjacent edges. Since the input represents a tree with n nodes, the total number of edges is n - 1, and each edge is examined exactly twice (once from each endpoint) across the entire traversal. Building the adjacency list g also takes O(n) time, as it processes all n - 1 edges. The final pow(2, d - 1, 10**9 + 7) operation runs in O(log d) time, which is dominated by the linear traversal. Therefore, the overall time complexity is O(n), where n is the number of nodes in the tree.

Space Complexity: O(n)

The adjacency list g stores all the edges of the tree, requiring O(n) space to hold the 2 * (n - 1) entries. Additionally, the dfs function is recursive, and in the worst case (e.g., a tree shaped like a linked list), the recursion depth can reach O(n), consuming O(n) stack space. Combining both factors, the total space complexity is O(n), where n is the number of nodes in the tree.

Pattern Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Off-by-One Error in the BFS Exponent

The most subtle and dangerous pitfall lies in the BFS alternative, where depth and edge count are confused.

In the BFS approach, depth is incremented once per level, including the level containing only the root. This means after the loop, depth equals the number of levels (nodes along the longest path), not the number of edges. Since edge count is levels - 1, the exponent must be depth - 2, not depth - 1:

# BFS: depth = number of levels (node count), so edges = depth - 1
# Odd-sum assignments = 2^(edges - 1) = 2^(depth - 2)
return pow(2, depth - 2, MOD)

Compare this to the recursive DFS, where dfs returns the edge count directly, so the exponent is correctly depth - 1:

# DFS: depth = edge count, so answer = 2^(depth - 1)
return pow(2, depth - 1, MOD)

Why it's easy to get wrong: Both versions use a variable named depth, but they measure different things (edges vs. levels). Blindly copying the exponent from one version into the other produces a result that is off by a factor of 2.

Solution: Always anchor your reasoning to the number of edges d. Pick one definition and verify it:

DFS returning edge count → pow(2, d - 1, MOD)
BFS counting levels → edges = levels - 1 → pow(2, levels - 2, MOD)

Pitfall 2: Negative or Zero Exponent When the Tree Is a Single Node

If n == 1, there are no edges, so edges is empty and the maximum depth d is 0. Computing pow(2, d - 1, MOD) becomes pow(2, -1, MOD).

In Python 3.8+, pow(2, -1, MOD) does not raise an error — it computes the modular inverse of 2, returning 500000004, which is mathematically meaningless for this problem. In older Python versions, it raises a ValueError.

Why it's dangerous: The code may silently return a wrong answer (the modular inverse) instead of crashing, making the bug very hard to detect.

Solution: Handle the empty-path case explicitly. With zero edges, there are zero ways to make an odd sum:

depth = dfs(1)
if depth == 0:
    return 0  # No edges → no way to form an odd path sum
return pow(2, depth - 1, MOD)

(Note: For this specific LeetCode problem, constraints typically guarantee n >= 2, but defensive handling avoids surprises if reused elsewhere.)

Pitfall 3: Recursion Depth Limit on Skewed (Chain-Like) Trees

The recursive dfs consumes one stack frame per level. For a degenerate tree shaped like a long chain (e.g., 1-2-3-...-n), the recursion depth equals n. With n in the range of 10^4–10^5, Python's default recursion limit (1000) is exceeded, raising a RecursionError.

Solution: Either raise the recursion limit:

import sys
sys.setrecursionlimit(10**6)

or use the iterative BFS version (which has no recursion at all), keeping the corrected depth - 2 exponent from Pitfall 1.

Pitfall 4: Incorrect Adjacency List Sizing

The adjacency list is allocated with size num_nodes + 1 to support 1-indexed labels (1 to n). A common mistake is allocating num_nodes instead, causing an IndexError when accessing graph[n]:

graph = [[] for _ in range(num_nodes)]      # ❌ IndexError at graph[n]
graph = [[] for _ in range(num_nodes + 1)]  # ✅ Correct for 1-indexed nodes

Solution: Always size 1-indexed structures as n + 1, and remember that index 0 simply goes unused.

Ready to land your dream job?

Unlock your dream job with a 5-minute quiz for a personalized study roadmap!

Get My Roadmap

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Get a Personalized Study Roadmap:

Which of the following is a good use case for backtracking?

Combinatorial problems

Random number generation

Topological sorting

Matrix multiplication

I don't know

Core Learning

LLD/OOP Interview

Live Training

Additional Resources

AI Engineering