LeetCode 3559 - Number of Ways to Assign Edge Weights II
We are given an undirected tree with n nodes, rooted at node 1. Every edge will eventually receive a weight of either 1 or 2. For each query [u, v], we only care about the unique path connecting u and v. All edges outside that path are completely irrelevant and should be ignored.
Difficulty: 🔴 Hard
Topics: Array, Math, Dynamic Programming, Bit Manipulation, Tree, Depth-First Search
Solution
Problem Understanding
We are given an undirected tree with n nodes, rooted at node 1. Every edge will eventually receive a weight of either 1 or 2.
For each query [u, v], we only care about the unique path connecting u and v. All edges outside that path are completely irrelevant and should be ignored.
The task is to count how many assignments of weights to the edges on that path produce an odd total path cost, where the path cost is the sum of all edge weights along the path.
The answer for every query must be returned modulo 10^9 + 7.
The most important observation is that the actual structure of the tree only determines how many edges are on the path. Once we know the path length, the counting problem becomes purely combinatorial.
The constraints are large:
n ≤ 100000queries.length ≤ 100000
This immediately rules out any approach that processes an entire path for every query. We need a way to answer path length queries efficiently.
Several edge cases are important:
- A query of the form
[u, u]has path length0. The cost is always0, so the answer is0. - A path containing exactly one edge has two possible assignments
{1,2}, and only one of them produces an odd sum. - Long paths can have lengths up to
n - 1, so powers of two must be computed modulo10^9 + 7. - Since there are up to
100000queries, path lengths must be computed in approximately logarithmic time. Let's carefully work through this LeetCode problem and provide a comprehensive technical guide following your formatting rules.
Problem Understanding
The problem gives us an undirected tree with n nodes labeled from 1 to n and n - 1 edges. Each edge can be assigned a weight of either 1 or 2. We are asked to answer queries about the number of ways to assign weights to edges in a path between two nodes such that the total weight along that path is odd. The output must be modulo 10^9 + 7.
In other words, for each query [u, v], we need to consider only the edges on the unique path from u to v and count how many combinations of assigning 1 or 2 to these edges result in an odd sum. Because the tree is connected and acyclic, there is exactly one path between any two nodes.
Constraints tell us that n can be up to 100,000, and there can be up to 100,000 queries. This is too large for any solution that tries to enumerate all assignments or walks through paths in linear time per query, so we need a clever mathematical approach.
Edge cases to note:
- A path from a node to itself has length 0, so the sum is 0, which is even. Valid assignments = 0.
- A path with one edge has exactly one valid assignment for an odd sum (assign 1 to that edge).
Approaches
Brute Force
For each query, we could explicitly find the path between the two nodes.
If the path contains k edges, there are 2^k possible weight assignments because every edge independently chooses between weight 1 and weight 2.
We could enumerate all assignments, compute the resulting path cost, and count those with odd sums.
This approach is correct because it checks every possible assignment directly. However, it is completely infeasible. A path may contain up to 100000 edges, making 2^k astronomically large.
Even if we only traversed the path and used dynamic programming per query, doing path work repeatedly for 100000 queries would still be too expensive.
Key Insight
The only odd weight is 1.
The only even weight is 2.
Therefore, the parity of the path sum depends solely on how many edges receive weight 1.
Let the path contain k edges.
The sum is odd if and only if the number of edges assigned weight 1 is odd.
Among all 2^k assignments, exactly half contain an odd number of chosen edges and half contain an even number.
Therefore:
- If
k = 0, answer =0 - If
k ≥ 1, answer =2^(k-1)
The problem now reduces to finding the number of edges on the path between two nodes.
For a rooted tree:
$$\text{distance}(u,v) = \text{depth}(u) + \text{depth}(v) - 2 \cdot \text{depth}(\text{LCA}(u,v))$$
Thus we only need efficient Lowest Common Ancestor queries.
Binary lifting computes LCA in O(log n) per query after O(n log n) preprocessing.
Approach Comparison
| Approach | Time Complexity | Space Complexity | Notes |
|---|---|---|---|
| Brute Force | O(2^k) per query |
O(k) |
Enumerates all weight assignments on the path |
| Optimal | O((n + q) log n) |
O(n log n) |
Use LCA to get path length, answer is 2^(length-1) |
Algorithm Walkthrough
Step 1: Build the tree
Convert the edge list into an adjacency list.
Since the graph is a tree, this representation allows efficient DFS traversal.
Step 2: Compute depths and binary lifting table
Run a DFS starting from root node 1.
For every node:
- Store its depth.
- Store its immediate parent.
- Fill binary lifting ancestors.
Let:
up[j][v]
represent the 2^j-th ancestor of node v.
This preprocessing allows jumping upward in logarithmic time.
Step 3: Precompute powers of two
The answer for a path of length k is:
$$2^{k-1}$$
Precompute:
$$pow2[i] = 2^i \pmod{10^9+7}$$
for all possible path lengths.
Step 4: Answer each query
For nodes u and v:
- Compute
lca = LCA(u, v). - Compute path length:
$$d = depth[u] + depth[v] - 2 \cdot depth[lca]$$
- If
d == 0, answer is0. - Otherwise answer is:
$$pow2[d-1]$$
Step 5: Return all answers
Store every query result in an array and return it.
Why it works
For a path containing d edges, each edge independently chooses weight 1 or 2.
Viewing weight 1 as parity 1 and weight 2 as parity 0, the path sum is odd exactly when an odd number of edges choose weight 1.
For any non-empty set of d binary choices, exactly half of the 2^d assignments have odd parity and half have even parity. Therefore the number of valid assignments is 2^(d-1).
The LCA formula correctly computes the number of edges on the path between two nodes. Combining these two facts yields the correct answer for every query. The brute-force approach would be to:
- For each query
[u, v], find the path betweenuandvin the tree. - Enumerate all
2^kcombinations of weights for thekedges in the path. - Count how many combinations give an odd sum.
This is correct but too slow because the path length k can be O(n), and with 2^k combinations per query, it is infeasible for the given constraints.
Key Insight
The key insight is that we do not need to enumerate all combinations. For a path with k edges:
- Each edge has two options: 1 or 2.
- The sum modulo 2 only depends on how many edges are assigned 1 (odd) and 2 (even).
- We want an odd sum, so the number of edges assigned 1 must be odd.
From combinatorics:
- Number of ways to pick an odd number of edges from
kedges to assign 1 =2^(k-1).
This works because the total number of subsets of k elements is 2^k, and exactly half of them have an odd number of elements. Therefore, for k edges:
- If
k = 0, answer = 0. - If
k >= 1, answer =2^(k-1).
To efficiently compute k, the distance between nodes in the tree is required. Using Depth-First Search (DFS) to precompute depths of all nodes relative to the root, the path length k between nodes u and v is:
k = depth[u] + depth[v] - 2 * depth[lca(u, v)]
Where lca(u, v) is the lowest common ancestor of u and v. Efficient LCA computation can be done in O(log n) per query using binary lifting.
| Approach | Time Complexity | Space Complexity | Notes |
|---|---|---|---|
| Brute Force | O(Q * 2^k) | O(n) | Enumerates all weight assignments along paths; infeasible |
| Optimal | O((n + Q) * log n) | O(n * log n) | Precompute LCA with binary lifting, then compute path length and 2^(k-1) per query |
Algorithm Walkthrough
-
Build the tree as an adjacency list from
edges. -
Run DFS from the root node (1) to record the depth and the parent of each node.
-
Precompute binary lifting table
up[node][i]for all nodes, whereup[node][i]is the 2^i-th ancestor of node. -
Implement a function to compute
lca(u, v)using the binary lifting table in O(log n) time. -
For each query
[u, v]: -
Compute the path length
kusingdepth[u] + depth[v] - 2 * depth[lca(u,v)]. -
If
k == 0, the answer is 0. -
Otherwise, compute
(2^(k-1)) % MODwhereMOD = 10^9 + 7. -
Collect all answers in an array and return.
Why it works:
The combinatorial property guarantees that for k edges, exactly half of the weight assignments yield an odd sum. Using LCA to compute the path length ensures correctness, and modulo operations handle large numbers safely.
Python Solution
from typing import List
import math
class Solution:
def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
MOD = 10**9 + 7
n = len(edges) + 1
graph = [[] for _ in range(n + 1)]
for u, v in edges:
graph[u].append(v)
graph[v].append(u)
LOG = math.ceil(math.log2(n)) + 1
up = [[0] * (n + 1) for _ in range(LOG)]
depth = [0] * (n + 1)
stack = [(1, 0)]
while stack:
node, parent = stack.pop()
up[0][node] = parent
for nei in graph[node]:
if nei == parent:
continue
depth[nei] = depth[node] + 1
stack.append((nei, node))
for j in range(1, LOG):
for v in range(1, n + 1):
up[j][v] = up[j - 1][up[j - 1][v]]
def lca(a: int, b: int) -> int:
if depth[a] < depth[b]:
a, b = b, a
diff = depth[a] - depth[b]
for j in range(LOG):
if diff & (1 << j):
a = up[j][a]
if a == b:
return a
for j in range(LOG - 1, -1, -1):
if up[j][a] != up[j][b]:
a = up[j][a]
b = up[j][b]
return up[0][a]
max_depth = max(depth)
pow2 = [1] * (max_depth + 1)
for i in range(1, len(pow2)):
pow2[i] = (pow2[i - 1] * 2) % MOD
answer = []
for u, v in queries:
ancestor = lca(u, v)
dist = (
depth[u]
+ depth[v]
- 2 * depth[ancestor]
)
if dist == 0:
answer.append(0)
else:
answer.append(pow2[dist - 1])
return answer
The implementation begins by constructing an adjacency list representation of the tree. A DFS from node 1 computes the depth of every node and records each node's direct parent.
After that, the binary lifting table is filled. Each entry stores a power-of-two ancestor, allowing nodes to be lifted upward efficiently during LCA computation.
The lca() function first equalizes depths and then lifts both nodes upward until their ancestors diverge. The parent immediately above that divergence is the LCA.
Once the LCA is known, the path length is computed using the standard distance formula. If the path length is zero, the answer is zero. Otherwise, the answer is simply the precomputed value 2^(distance-1) modulo 10^9 + 7.
MOD = 10**9 + 7
class Solution: def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]: from collections import defaultdict, deque
n = len(edges) + 1
tree = defaultdict(list)
for u, v in edges:
tree[u].append(v)
tree[v].append(u)
LOG = 17 # 2^17 > 1e5
parent = [[-1] * (LOG + 1) for _ in range(n + 1)]
depth = [0] * (n + 1)
def dfs(u, p):
parent[u][0] = p
for v in tree[u]:
if v != p:
depth[v] = depth[u] + 1
dfs(v, u)
dfs(1, -1)
for j in range(1, LOG + 1):
for i in range(1, n + 1):
if parent[i][j - 1] != -1:
parent[i][j] = parent[parent[i][j - 1]][j - 1]
def lca(u, v):
if depth[u] < depth[v]:
u, v = v, u
for i in range(LOG, -1, -1):
if parent[u][i] != -1 and depth[parent[u][i]] >= depth[v]:
u = parent[u][i]
if u == v:
return u
for i in range(LOG, -1, -1):
if parent[u][i] != -1 and parent[u][i] != parent[v][i]:
u = parent[u][i]
v = parent[v][i]
return parent[u][0]
def mod_pow2(x):
return pow(2, x, MOD)
res = []
for u, v in queries:
k = depth[u] + depth[v] - 2 * depth[lca(u, v)]
if k == 0:
res.append(0)
else:
res.append(mod_pow2(k - 1))
return res
**Explanation:**
- DFS builds depth and immediate parent relationships.
- Binary lifting allows O(log n) LCA queries.
- For each query, path length is computed and used in the combinatorial formula `2^(k-1)` modulo `10^9 + 7`.
- Edge cases like `u == v` (no edges) are handled.
## Go Solution
```go
package main
import "math/bits"
func assignEdgeWeights(edges [][]int, queries [][]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)
}
LOG := bits.Len(uint(n)) + 1
up := make([][]int, LOG)
for i := range up {
up[i] = make([]int, n+1)
}
depth := make([]int, n+1)
type Pair struct {
node int
parent int
}
stack := []Pair{{1, 0}}
for len(stack) > 0 {
cur := stack[len(stack)-1]
stack = stack[:len(stack)-1]
node := cur.node
parent := cur.parent
up[0][node] = parent
for _, nei := range graph[node] {
if nei == parent {
continue
}
depth[nei] = depth[node] + 1
stack = append(stack, Pair{nei, node})
}
}
for j := 1; j < LOG; j++ {
for v := 1; v <= n; v++ {
up[j][v] = up[j-1][up[j-1][v]]
}
}
lca := func(a, b int) int {
if depth[a] < depth[b] {
a, b = b, a
}
diff := depth[a] - depth[b]
for j := 0; j < LOG; j++ {
if diff&(1<<j) != 0 {
a = up[j][a]
}
}
if a == b {
return a
}
for j := LOG - 1; j >= 0; j-- {
if up[j][a] != up[j][b] {
a = up[j][a]
b = up[j][b]
}
}
return up[0][a]
}
maxDepth := 0
for i := 1; i <= n; i++ {
if depth[i] > maxDepth {
maxDepth = depth[i]
}
}
pow2 := make([]int, maxDepth+1)
pow2[0] = 1
for i := 1; i <= maxDepth; i++ {
pow2[i] = int((int64(pow2[i-1]) * 2) % int64(MOD))
}
ans := make([]int, 0, len(queries))
for _, q := range queries {
u, v := q[0], q[1]
ancestor := lca(u, v)
dist := depth[u] + depth[v] - 2*depth[ancestor]
if dist == 0 {
ans = append(ans, 0)
} else {
ans = append(ans, pow2[dist-1])
}
}
return ans
}
The Go implementation follows exactly the same algorithm. The primary differences are the explicit slice allocation for the binary lifting table and the use of int64 during modular multiplication to avoid overflow. An iterative DFS is used, which avoids recursion depth issues on large trees.
Worked Examples
Example 1
Input:
edges = [[1,2]]
queries = [[1,1],[1,2]]
Tree:
1
|
2
Depth table:
| Node | Depth |
|---|---|
| 1 | 0 |
| 2 | 1 |
Query [1,1]
| Value | Result |
|---|---|
| LCA(1,1) | 1 |
| Distance | 0 + 0 - 2×0 = 0 |
Since distance is 0, there are no edges.
Answer:
0
Query [1,2]
| Value | Result |
|---|---|
| LCA(1,2) | 1 |
| Distance | 0 + 1 - 2×0 = 1 |
Number of assignments:
$$2^{1-1}=1$$
Answer:
1
Final output:
[0,1]
Example 2
Input:
edges = [[1,2],[1,3],[3,4],[3,5]]
queries = [[1,4],[3,4],[2,5]]
Depth table:
| Node | Depth |
|---|---|
| 1 | 0 |
| 2 | 1 |
| 3 | 1 |
| 4 | 2 |
| 5 | 2 |
Query [1,4]
| Value | Result |
|---|---|
| LCA | 1 |
| Distance | 0 + 2 - 0 = 2 |
Answer:
$$2^{2-1}=2$$
Query [3,4]
| Value | Result |
|---|---|
| LCA | 3 |
| Distance | 1 + 2 - 2 = 1 |
Answer:
$$2^{1-1}=1$$
Query [2,5]
| Value | Result |
|---|---|
| LCA | 1 |
| Distance | 1 + 2 - 0 = 3 |
Answer:
$$2^{3-1}=4$$
Final output:
[2,1,4]
Complexity Analysis
| Measure | Complexity | Explanation |
|---|---|---|
| Time | O((n + q) log n) |
Binary lifting preprocessing plus one LCA query per request |
| Space | O(n log n) |
Ancestor table and tree storage |
The adjacency list requires O(n) space. The binary lifting table stores log n ancestors for every node, giving O(n log n) space. Each query performs an LCA computation in O(log n) time, resulting in total query cost O(q log n).
Test Cases
from typing import List
s = Solution()
# Example 1
assert s.assignEdgeWeights(
[[1, 2]],
[[1, 1], [1, 2]]
) == [0, 1]
# Example 2
assert s.assignEdgeWeights(
[[1, 2], [1, 3], [3, 4], [3, 5]],
[[1, 4], [3, 4], [2, 5]]
) == [2, 1, 4]
# Single edge queried in reverse direction
assert s.assignEdgeWeights(
[[1, 2]],
[[2, 1]]
) == [1]
# Chain of length 3
assert s.assignEdgeWeights(
[[1, 2], [2, 3], [3, 4]],
[[1, 4]]
) == [4] # distance=3 => 2^(3-1)=4
# Same node queries
assert s.assignEdgeWeights(
[[1, 2], [2, 3]],
[[1, 1], [2, 2], [3, 3]]
) == [0, 0, 0]
# Star tree
assert s.assignEdgeWeights(
[[1, 2], [1, 3], [1, 4], [1, 5]],
[[2, 3], [4, 5]]
) == [2, 2] # distance=2
# Mixed queries
assert s.assignEdgeWeights(
[[1, 2], [2, 3], [2, 4], [4, 5]],
[[3, 5], [1, 5], [3, 4]]
) == [4, 4, 2]
# Long chain stress-style correctness
n = 10
edges = [[i, i + 1] for i in range(1, n)]
assert s.assignEdgeWeights(
edges,
[[1, 10]]
) == [256] # distance=9 => 2^8
Test Summary
| Test | Why |
|---|---|
| Example 1 | Verifies zero-length and one-edge paths |
| Example 2 | Verifies general tree structure |
| Reverse single edge | Confirms path direction does not matter |
| Chain length 3 | Checks larger distance formula |
| Same node queries | Validates distance zero handling |
| Star tree | Tests LCA at root |
| Mixed queries | Tests multiple LCA configurations |
| Long chain | Verifies power-of-two counting for larger distances |
Edge Cases
Querying a Node with Itself
A query such as [u, u] produces a path containing zero edges. A common mistake is to apply the formula 2^(d-1) blindly, which would attempt to compute 2^(-1).
The implementation explicitly checks whether dist == 0 and returns 0 immediately.
Path Length Equal to One
When two nodes are connected by a single edge, there are exactly two assignments: weight 1 or weight 2. Only weight 1 yields an odd sum.
The formula gives:
$$2^{1-1}=1$$
which correctly matches the expected count.
Very Deep Trees
A tree can contain up to 100000 nodes and may be a straight chain. Recursive DFS implementations can exceed recursion limits.
Both provided solutions use iterative traversal, avoiding stack overflow issues while still computing depths and ancestors correctly. const MOD int = 1e9 + 7
func assignEdgeWeights(edges [][]int, queries [][]int) []int { n := len(edges) + 1 tree := make([][]int, n+1) for _, e := range edges { u, v := e[0], e[1] tree[u] = append(tree[u], v) tree[v] = append(tree[v], u) }
LOG := 17
parent := make([][]int, n+1)
for i := range parent {
parent[i] = make([]int, LOG+1)
for j := range parent[i] {
parent[i][j] = -1
}
}
depth := make([]int, n+1)
var dfs func(int, int)
dfs = func(u, p int) {
parent[u][0] = p
for _, v := range tree[u] {
if v != p {
depth[v] = depth[u] + 1
dfs(v, u)
}
}
}
dfs(1, -1)
for j := 1; j <= LOG; j++ {
for i := 1; i <= n; i++ {
if parent[i][j-1] != -1 {
parent[i][j] = parent[parent[i][j-1]][j-1]
}
}
}
lca := func(u, v int) int {
if depth[u] < depth[v] {
u, v = v, u
}
for i := LOG; i >= 0;