LeetCode 3558 - Number of Ways to Assign Edge Weights I

We are given an undirected tree rooted at node 1. Every edge initially has weight 0, and we must assign each edge a weight of either 1 or 2. The problem asks us to choose any node x that has the maximum depth in the tree.

leetcodemediummathtreedepth-first-search

LeetCode Problem 3558

Difficulty: 🟡 Medium
Topics: Math, Tree, Depth-First Search

Solution

LeetCode 3558 - Number of Ways to Assign Edge Weights I

Problem Understanding

We are given an undirected tree rooted at node 1. Every edge initially has weight 0, and we must assign each edge a weight of either 1 or 2.

The problem asks us to choose any node x that has the maximum depth in the tree. We then consider only the unique path from the root node 1 to x. All edges outside this path are irrelevant and can be ignored completely.

Our goal is to count how many assignments of weights 1 and 2 to the edges on this root-to-x path produce an odd total path cost. Since the answer can be very large, we return it modulo 10^9 + 7.

A key detail is that if multiple nodes exist at the maximum depth, the answer is the same regardless of which one is chosen. The only thing that matters is the length of the root-to-deepest-node path, because every deepest node has the same depth.

The constraints are important:

  • n can be as large as 10^5.
  • The graph is guaranteed to be a valid tree.
  • A tree with n nodes contains exactly n - 1 edges.

These constraints immediately rule out any solution that tries all possible weight assignments, because a path can contain up to 10^5 - 1 edges.

Several edge cases are worth noting:

  • The smallest possible tree contains only one edge.
  • The tree may be a straight chain, creating a maximum depth of n - 1.
  • The tree may be highly branched, with many deepest nodes sharing the same depth.
  • Recursive DFS may overflow the call stack when the tree is a long chain, so an iterative traversal is safer.

Approaches

Brute Force

Suppose the root-to-deepest-node path contains d edges.

Each edge can receive either weight 1 or weight 2, so there are:

$$2^d$$

possible assignments.

For every assignment, we could compute the path sum and check whether it is odd. This approach is correct because it explicitly examines every possible configuration.

However, when d approaches 10^5, the number of assignments becomes astronomically large. Even for d = 50, 2^d is already far beyond practical limits.

Therefore, brute force is completely infeasible.

Key Insight

Only parity matters.

Notice:

  • Weight 1 contributes odd parity.
  • Weight 2 contributes even parity.

Therefore, the parity of the total path cost depends only on how many edges receive weight 1.

The sum is odd if and only if the number of edges assigned weight 1 is odd.

Now suppose the path contains d edges.

Every edge independently chooses between:

  • weight 1 (odd)
  • weight 2 (even)

Among all 2^d assignments:

  • Exactly half produce an odd sum.
  • Exactly half produce an even sum.

This follows from a standard parity argument. Flipping the first edge between 1 and 2 changes the parity of the total sum, creating a one-to-one correspondence between odd-sum and even-sum assignments.

Therefore, the answer is simply:

$$\frac{2^d}{2} = 2^{d-1}$$

for any d ≥ 1.

The remaining task is to find the maximum depth d of the tree.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(2^d · d) O(d) Enumerates all weight assignments
Optimal O(n) O(n) Find maximum depth, then compute $2^{d-1}$

Algorithm Walkthrough

  1. Let n = len(edges) + 1.
  2. Build an adjacency list for the tree so that all neighbors of every node can be accessed efficiently.
  3. Perform an iterative DFS (or BFS) starting from node 1.
  4. Track the depth of every visited node. Since the tree is rooted at node 1, the depth of a child is the depth of its parent plus one.
  5. Maintain a variable max_depth, updating it whenever a deeper node is discovered.
  6. After traversal completes, max_depth equals the number of edges on the path from the root to any deepest node.
  7. Let d = max_depth.
  8. The number of valid assignments is:

$$2^{d-1}$$

because exactly half of all 2^d assignments have odd parity.

  1. Return:

$$2^{d-1} \bmod (10^9+7)$$

using fast modular exponentiation.

Why it works

Every edge contributes either odd parity (1) or even parity (2). The total path cost is odd exactly when an odd number of edges receive weight 1.

For a path of length d, there are 2^d assignments. Toggling the weight of any fixed edge, for example the first edge, changes the parity of the total sum and creates a bijection between odd-sum assignments and even-sum assignments. Therefore exactly half the assignments are odd, giving 2^(d-1) valid assignments. Since all deepest nodes have the same depth, only the maximum depth matters.

Problem Understanding

We are given an undirected tree rooted at node 1, with n nodes and n - 1 edges. Every edge initially has weight 0, and we must assign each edge a weight of either 1 or 2. The assignment is only relevant for a single root-to-leaf path: we first identify any node x that lies at the maximum depth from the root, and then we consider only the unique path from node 1 to that node x.

The task is to count how many ways we can assign weights (1 or 2) to the edges along that path such that the total sum of weights is odd. The answer must be returned modulo 1e9 + 7.

The input is an edge list representing a valid tree, so there are no cycles and exactly one simple path between any two nodes. The constraints allow up to 10^5 nodes, which implies any exponential enumeration over assignments is infeasible except in a highly reduced form.

A key observation is that we do not need to consider all nodes or all paths. We only need the depth of the deepest node from the root, because every deepest node has the same depth, and any such node yields a path of the same length in terms of number of edges.

Important edge cases include a star-shaped tree where all nodes are at depth 1, and a skewed tree where depth is n - 1. Another subtle case is when multiple nodes share maximum depth, but the answer depends only on path length, not which node is chosen.

Approaches

Brute Force Approach

A brute force method would involve first finding a deepest node, then enumerating all possible assignments of weights 1 or 2 along the path from root to that node. For a path of length d, there are 2^d possible assignments. For each assignment, we compute the sum and check whether it is odd.

This approach is correct because it exhaustively evaluates every possible configuration, but it is far too slow. When d can be as large as 10^5, enumerating 2^d assignments is impossible.

Key Insight

The critical observation is that the structure of the tree becomes irrelevant once we reduce the problem to a single path. Each edge independently contributes either 1 (odd) or 2 (even). Therefore, the parity of the total sum depends only on how many edges are assigned weight 1.

Let the path length be d. We want the number of assignments where the sum is odd. Since only weight 1 contributes to parity, this becomes a combinatorics problem: count subsets of edges assigned value 1 such that the subset size is odd.

This is a classic identity:

The number of subsets of size odd in a set of size d is 2^(d-1) for d ≥ 1.

So the entire problem reduces to:

  1. Compute maximum depth d from root 1.
  2. Return 2^(d-1) mod (1e9 + 7).

Complexity Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(2^d · d) O(d) Enumerates all weight assignments on the path
Optimal O(n) O(n) DFS/BFS to compute depth, then modular exponentiation

Algorithm Walkthrough

The optimal solution proceeds as follows.

  1. First, construct an adjacency list representation of the tree from the edge list. This is necessary because efficient traversal requires quick access to neighbors of each node.
  2. Perform a depth-first search (DFS) or breadth-first search (BFS) starting from node 1. The goal is to compute the maximum depth of any node reachable from the root. We maintain a depth array where depth[u] represents the distance (in edges) from node 1 to node u.
  3. During traversal, for each neighbor of a node, if it has not been visited, we assign its depth as depth[current] + 1 and continue exploring.
  4. Track the maximum depth encountered during traversal. This value corresponds to the number of edges in the root-to-deepest-node path.
  5. Once traversal is complete, let the maximum depth be d. If d == 0, the answer is trivially 0 (though this case does not occur under constraints since n ≥ 2).
  6. Compute the result as 2^(d - 1) mod (1e9 + 7) using fast exponentiation.

Why it works

The correctness relies on two key properties. First, all deepest nodes have the same depth, so the number of edges in any root-to-deepest path is identical. Second, each edge independently contributes either an odd or even value, so the parity of the sum depends only on how many edges are assigned value 1. This reduces the problem to counting subsets of odd size among d edges, which is exactly 2^(d-1).

Python Solution

from typing import List
from collections import defaultdict
from collections import defaultdict, deque

class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        MOD = 10**9 + 7
        n = len(edges) + 1

        
        # Build adjacency list
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        max_depth = 0
        stack = [(1, 0, 0)]  # (node, parent, depth)

        while stack:
            node, parent, depth = stack.pop()
            max_depth = max(max_depth, depth)

            for neighbor in graph[node]:
                if neighbor != parent:
                    stack.append((neighbor, node, depth + 1))

        return pow(2, max_depth - 1, MOD)

Implementation Explanation

The adjacency list represents the tree in both directions because the input edges are undirected.

An iterative DFS starts from node 1. Each stack entry stores the current node, its parent, and its depth. The parent information prevents revisiting the node we came from.

During traversal, max_depth is updated whenever a deeper node is reached. When DFS finishes, max_depth is exactly the length of the path from the root to any deepest node.

Once the depth is known, the mathematical observation eliminates any need to examine weight assignments. The answer is simply 2^(max_depth - 1) modulo 10^9 + 7, computed efficiently using Python's built-in modular exponentiation.

    # BFS to compute depths from root 1
    n = len(edges) + 1
    depth = {1: 0}
    q = deque([1])
    visited = set([1])
    
    max_depth = 0
    
    while q:
        node = q.popleft()
        for nei in graph[node]:
            if nei not in visited:
                visited.add(nei)
                depth[nei] = depth[node] + 1
                max_depth = max(max_depth, depth[nei])
                q.append(nei)
    
    # If only root exists
    if max_depth == 0:
        return 0
    
    # number of valid assignments = 2^(d-1)
    return pow(2, max_depth - 1, MOD)


### Code Explanation

We first construct an adjacency list for efficient traversal of the tree. We then perform BFS starting from node 1, computing the shortest distance to every node, which in a tree corresponds to its depth from the root. During BFS, we track the maximum depth encountered. Finally, we compute the combinatorial result using modular exponentiation.

## Go Solution

```go
package main

func assignEdgeWeights(edges [][]int) int {
	const MOD int = 1_000_000_007

	n := len(edges) + 1

	graph := make([][]int, n+1)
	for _, e := range edges {
		u, v := e[0], e[1]
		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	type State struct {
		node   int
		parent int
		depth  int
	}

	stack := []State{{1, 0, 0}}
	maxDepth := 0

	for len(stack) > 0 {
		cur := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if cur.depth > maxDepth {
			maxDepth = cur.depth
		}

		for _, next := range graph[cur.node] {
			if next != cur.parent {
				stack = append(stack, State{
					node:   next,
					parent: cur.node,
					depth:  cur.depth + 1,
				})
			}
		}
	}

	return modPow(2, maxDepth-1, MOD)
}

func modPow(base, exp, mod int) int {
	result := 1
	b := base % mod

	for exp > 0 {
		if exp&1 == 1 {
			result = (result * b) % mod
		}
		b = (b * b) % mod
		exp >>= 1
	}

	return result
}

Go-Specific Notes

The Go implementation uses an explicit stack instead of recursion, avoiding stack overflow on very deep trees.

A custom modPow function performs binary exponentiation because Go does not provide a built-in modular power function like Python's pow(base, exp, mod).

Since the modulus is below 2^31, standard integer arithmetic is sufficient for this problem. func assignEdgeWeights(edges [][]int) int { const MOD int = 1000000007

n := len(edges) + 1

// Build adjacency list
graph := make(map[int][]int, n)
for _, e := range edges {
    u, v := e[0], e[1]
    graph[u] = append(graph[u], v)
    graph[v] = append(graph[v], u)
}

// BFS
type nodeDepth struct {
    node  int
    depth int
}

visited := make(map[int]bool)
queue := []nodeDepth{{1, 0}}
visited[1] = true

maxDepth := 0
idx := 0

for idx < len(queue) {
    cur := queue[idx]
    idx++

    for _, nei := range graph[cur.node] {
        if !visited[nei] {
            visited[nei] = true
            nd := nodeDepth{nei, cur.depth + 1}
            queue = append(queue, nd)
            if nd.depth > maxDepth {
                maxDepth = nd.depth
            }
        }
    }
}

if maxDepth == 0 {
    return 0
}

// compute 2^(maxDepth-1) mod MOD
res := 1
base := 2
exp := maxDepth - 1

for exp > 0 {
    if exp%2 == 1 {
        res = (res * base) % MOD
    }
    base = (base * base) % MOD
    exp /= 2
}

return res

}


### Go-specific Notes

The Go implementation uses a slice-based BFS queue instead of a deque. We explicitly track indices to avoid costly popping from the front. Modular exponentiation is implemented manually since Go lacks a built-in `pow mod` function for integers.

## Worked Examples

### Example 1

Input:

edges = [[1,2]]


Tree:

1 | 2


DFS traversal:

| Node | Depth | max_depth |
| --- | --- | --- |
| 1 | 0 | 0 |
| 2 | 1 | 1 |

After traversal:

max_depth = 1


Answer:

$$2^{1-1}=2^0=1$$

Output:

1


### Example 2

Input:

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


Tree:


1

/
2 3 /
4 5


DFS traversal:

| Node | Depth | max_depth |
| --- | --- | --- |
| 1 | 0 | 0 |
| 2 | 1 | 1 |
| 3 | 1 | 1 |
| 4 | 2 | 2 |
| 5 | 2 | 2 |

After traversal:

max_depth = 2


The root-to-deepest-node path contains two edges.

Total assignments:

$$2^2 = 4$$

They are:

| Assignment | Sum | Odd? |
| --- | --- | --- |
| (1,1) | 2 | No |
| (1,2) | 3 | Yes |
| (2,1) | 3 | Yes |
| (2,2) | 4 | No |

Valid assignments:

2


Formula:

$$2^{2-1}=2$$

Output:

2

Input: `edges = [[1,2]]`

Step 1: Build graph

Node 1 connects to node 2.

Step 2: BFS traversal

- depth(1) = 0
- depth(2) = 1

Maximum depth = 1

Step 3: Compute result

Path length `d = 1`

Result = `2^(1-1) = 2^0 = 1`

Only one edge, so only assignments `{1}` or `{2}` exist, and only `{1}` yields odd sum.

### Example 2

Input: `edges = [[1,2],[1,3],[3,4],[3,5]]`

Step 1: Build graph

1 connects to 2 and 3, 3 connects to 4 and 5.

Step 2: BFS traversal

- depth(1) = 0
- depth(2) = 1
- depth(3) = 1
- depth(4) = 2
- depth(5) = 2

Maximum depth = 2

Step 3: Compute result

Path length `d = 2`

Result = `2^(2-1) = 2`

Valid assignments are those with exactly one edge assigned weight 1.

## Complexity Analysis

| Measure | Complexity | Explanation |
| --- | --- | --- |
| Time | O(n) | Single traversal of all nodes and edges |
| Space | O(n) | Adjacency list and traversal stack |

The tree contains exactly `n - 1` edges. Building the adjacency list takes `O(n)` time. The DFS visits each node once and each edge twice, resulting in overall linear complexity. The adjacency list and stack together require linear auxiliary space.
| Time | O(n) | BFS visits each node and edge once, plus O(log d) exponentiation |
| Space | O(n) | Adjacency list and visitation tracking for all nodes |

The overall complexity is linear in the size of the tree because each edge is processed exactly twice during BFS construction and traversal.

## Test Cases

```python
from typing import List

s = Solution()

assert s.assignEdgeWeights([[1, 2]]) == 1  # minimum tree

assert s.assignEdgeWeights([
    [1, 2],
    [1, 3],
    [3, 4],
    [3, 5]
]) == 2  # provided example

assert s.assignEdgeWeights([
    [1, 2],
    [2, 3]
]) == 2  # depth 2, answer = 2^(2-1)

assert s.assignEdgeWeights([
    [1, 2],
    [2, 3],
    [3, 4]
]) == 4  # depth 3, answer = 2^(3-1)

assert s.assignEdgeWeights([
    [1, 2],
    [1, 3],
    [1, 4],
    [1, 5]
]) == 1  # star tree, max depth = 1

assert s.assignEdgeWeights([
    [1, 2],
    [2, 3],
    [2, 4],
    [4, 5]
]) == 4  # deepest depth = 3

assert s.assignEdgeWeights([
    [1, 2],
    [2, 3],
    [3, 4],
    [4, 5],
    [5, 6]
]) == 16  # chain of depth 5

assert s.assignEdgeWeights([
    [1, 2],
    [1, 3],
    [2, 4],
    [2, 5],
    [3, 6],
    [3, 7]
]) == 2  # balanced tree, depth 2

Test Summary

Test Why
[[1,2]] Smallest valid tree
Example 2 tree Verifies sample output
Chain of length 2 Checks depth calculation
Chain of length 3 Verifies formula for larger depth
Star tree Many deepest nodes at depth 1
Unbalanced tree Deepest node not on first branch
Long chain Maximum-depth style structure
Balanced binary tree Multiple deepest nodes with same depth

Edge Cases

Minimum Tree Size

The smallest valid input contains exactly two nodes and one edge. The maximum depth is 1, so the formula becomes 2^(1-1) = 1. This case is important because exponent zero must be handled correctly. The implementation naturally handles it through modular exponentiation.

Multiple Deepest Nodes

A tree may have several nodes at the maximum depth. At first glance it may seem necessary to determine which deepest node to choose. However, all deepest nodes have the same depth, and the answer depends only on the path length. The DFS therefore only tracks the maximum depth and does not need to identify a specific deepest node.

Deep Linear Tree

A tree shaped like a linked list can have depth close to 10^5. Recursive DFS implementations may overflow the call stack in such cases. The solution uses an explicit stack for iterative DFS, ensuring it remains safe and efficient even at the maximum constraint size.

Highly Branched Tree

A star-shaped tree can have many children directly connected to the root. The maximum depth is still only 1. The traversal correctly computes this depth regardless of the branching factor, and the formula returns 1, which matches the single valid odd assignment for a one-edge path. class Solution: def assignEdgeWeights(self, edges: List[List[int]]) -> int: MOD = 10**9 + 7 from collections import defaultdict, deque

    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)

    depth = {1: 0}
    q = deque([1])
    visited = {1}

    max_depth = 0

    while q:
        node = q.popleft()
        for nei in graph[node]:
            if nei not in visited:
                visited.add(nei)
                depth[nei] = depth[node] + 1
                max_depth = max(max_depth, depth[nei])
                q.append(nei)

    if max_depth == 0:
        return 0

    return pow(2, max_depth - 1, MOD)

basic example

assert Solution().assignEdgeWeights([[1,2]]) == 1 # single edge

example 2

assert Solution().assignEdgeWeights([[1,2],[1,3],[3,4],[3,5]]) == 2

skewed tree

assert Solution().assignEdgeWeights([[1,2],[2,3],[3,4],[4,5]]) == 8 # depth=4 => 2^3

star tree

assert Solution().assignEdgeWeights([[1,2],[1,3],[1,4],[1,5]]) == 1 # depth=1

larger balanced tree

assert Solution().assignEdgeWeights([[1,2],[1,3],[2,6],[2,7],[3,8],[3,9]]) == 2


| Test | Why |
| --- | --- |
| single edge | minimal non-trivial tree |
| branching tree | verifies multiple deepest nodes |
| skewed tree | tests maximum depth growth |
| star tree | tests shallow depth case |
| balanced tree | verifies correct depth selection |

## Edge Cases

One important edge case is a star-shaped tree where node 1 is connected directly to all other nodes. In this case, the maximum depth is 1, meaning the path has only one edge. The formula correctly returns `2^(1-1) = 1`, avoiding incorrect assumptions about multiple levels.

Another edge case is a completely skewed tree where every node forms a chain. This maximizes depth to `n - 1`. The implementation handles this naturally through BFS without stack overflow risks because it uses iterative traversal.

A third case is when multiple nodes share the maximum depth. This does not affect correctness because the answer depends only on the depth value, not the identity of the node. The BFS tracks only the maximum depth, ensuring all such nodes are implicitly covered.