Number of Ways to Assign Edge Weights I

Med

#3168Time: O(n + 2^L). Finding the max depth `L` takes `O(n)` time using BFS. The recursive function has two branches at each of the `L` levels, leading to `O(2^L)` calls. This is exponential and too slow for the given constraints.Space: O(n + L). `O(n)` for the adjacency list and BFS data structures. The recursion depth can go up to `L`, so the call stack uses `O(L)` space. Since `L < n`, this is `O(n)`.

Patterns

Math

Algorithms

Depth-First Search

Data structures

Tree

Prompt

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_i, v_i] indicates that there is an edge between nodes u_i and v_i.

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.

Example 1:

Input: edges = [[1,2]]

Output: 1

Explanation:

The path from Node 1 to Node 2 consists of one edge (1 → 2).
Assigning weight 1 makes the cost odd, while 2 makes it even. Thus, the number of valid assignments is 1.

Example 2:

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

Output: 2

Explanation:

The maximum depth is 2, with nodes 4 and 5 at the same depth. Either node can be selected for processing.
For example, the path from Node 1 to Node 4 consists of two edges (1 → 3 and 3 → 4).
Assigning weights (1,2) or (2,1) results in an odd cost. Thus, the number of valid assignments is 2.

Constraints:

2 <= n <= 10⁵
edges.length == n - 1
edges[i] == [u_i, v_i]
1 <= u_i, v_i <= n
edges represents a valid tree.

Approaches

3 approaches with complexity analysis and trade-offs.

This approach first determines the maximum depth of the tree, then explicitly generates every possible weight assignment for a path of that length using recursion. It checks each assignment's total cost and counts how many are odd. This method directly simulates the problem statement but is computationally very expensive.

Algorithm

Build Graph: Construct an adjacency list representation of the tree from the edges array.
Find Max Depth (L): Perform a Breadth-First Search (BFS) or Depth-First Search (DFS) starting from the root (node 1) to find the maximum depth L of any node in the tree. The path from the root to such a node will have L edges.
Recursive Enumeration: Create a recursive function, say countWays(edgeIndex, currentParity), where edgeIndex is the current edge on the path (from 0 to L-1) and currentParity is the parity of the sum so far (0 for even, 1 for odd).
- The base case is when edgeIndex == L. If currentParity is 1 (odd), return 1; otherwise, return 0.
- In the recursive step, for the edgeIndex-th edge, explore both possibilities:
  - Assign weight 1 (odd): This flips the parity. Recursively call countWays(edgeIndex + 1, 1 - currentParity).
  - Assign weight 2 (even): This preserves the parity. Recursively call countWays(edgeIndex + 1, currentParity).
- The total ways is the sum of the results from these two recursive calls (modulo 10^9 + 7).
Initial Call: Start the recursion with countWays(0, 0) (0 edges processed, sum is 0 which is even).
Return Result: The final count is the answer.

Walkthrough

The core idea is to solve the problem by exploring all possibilities. First, we need to know the number of edges in the path, which is simply the maximum depth of the tree, let's call it L. We can find L using a standard graph traversal like BFS starting from the root (node 1). Once L is known, we can define a recursive function that tries assigning weight 1 and weight 2 to each of the L edges. For each of the 2^L total combinations, we calculate the sum and check its parity. A counter is incremented for each combination that results in an odd sum.

1class Solution {2    private long MOD = 1000000007;3 4    public int numberOfWays(int[][] edges) {5        int n = edges.length + 1;6        if (n <= 1) return 0;7 8        List<List<Integer>> adj = new ArrayList<>();9        for (int i = 0; i <= n; i++) {10            adj.add(new ArrayList<>());11        }12        for (int[] edge : edges) {13            adj.get(edge[0]).add(edge[1]);14            adj.get(edge[1]).add(edge[0]);15        }16 17        int maxDepth = findMaxDepth(n, adj);18        if (maxDepth == 0) return 0;19 20        return (int) countRecursive(0, maxDepth, 0);21    }22 23    private long countRecursive(int edgeIndex, int L, int currentParity) {24        if (edgeIndex == L) {25            return currentParity; // Returns 1 if odd, 0 if even26        }27 28        // Case 1: Assign weight 1 (odd)29        long ways1 = countRecursive(edgeIndex + 1, L, 1 - currentParity);30 31        // Case 2: Assign weight 2 (even)32        long ways2 = countRecursive(edgeIndex + 1, L, currentParity);33 34        return (ways1 + ways2) % MOD;35    }36 37    private int findMaxDepth(int n, List<List<Integer>> adj) {38        Queue<int[]> queue = new LinkedList<>();39        queue.offer(new int[]{1, 0}); // {node, depth}40        boolean[] visited = new boolean[n + 1];41        visited[1] = true;42        int maxDepth = 0;43 44        while (!queue.isEmpty()) {45            int[] current = queue.poll();46            int u = current[0];47            int depth = current[1];48            maxDepth = Math.max(maxDepth, depth);49 50            for (int v : adj.get(u)) {51                if (!visited[v]) {52                    visited[v] = true;53                    queue.offer(new int[]{v, depth + 1});54                }55            }56        }57        return maxDepth;58    }59}

Complexity

Time

O(n + 2^L). Finding the max depth `L` takes `O(n)` time using BFS. The recursive function has two branches at each of the `L` levels, leading to `O(2^L)` calls. This is exponential and too slow for the given constraints.

Space

O(n + L). `O(n)` for the adjacency list and BFS data structures. The recursion depth can go up to `L`, so the call stack uses `O(L)` space. Since `L < n`, this is `O(n)`.

Trade-offs

Pros

Simple to understand and follows the problem definition directly.

Cons

Extremely inefficient due to its exponential time complexity.
Will cause a 'Time Limit Exceeded' error on platforms like LeetCode for the given constraints.

This approach also starts by finding the maximum depth L. Instead of brute-force, it uses dynamic programming to efficiently count the number of ways. We build up the solution for a path of length i based on the solutions for a path of length i-1, keeping track of the number of ways to form even and odd sums.

Algorithm

Build Graph & Find Max Depth (L): As in the previous approach, use BFS to find the maximum depth L.
DP State: Define a DP state based on the number of edges processed. Let evenWays[i] be the number of ways to get an even sum using i edges, and oddWays[i] be the number of ways to get an odd sum.
DP Initialization: For a single edge (i=1):
- evenWays[1] = 1 (by assigning weight 2).
- oddWays[1] = 1 (by assigning weight 1).
DP Recurrence: For i from 2 to L:
- To get an even sum for i edges: (even sum from i-1 + even weight) OR (odd sum from i-1 + odd weight). So, evenWays[i] = evenWays[i-1] + oddWays[i-1].
- To get an odd sum for i edges: (even sum from i-1 + odd weight) OR (odd sum from i-1 + even weight). So, oddWays[i] = evenWays[i-1] + oddWays[i-1].
Space Optimization: Notice that the calculation for i only depends on i-1. We can use two variables instead of arrays to store the counts for the previous step, reducing space to O(1) for the DP part.
Result: The answer is oddWays[L] after iterating up to L.

Walkthrough

After finding the maximum path length L, we can solve the counting problem using dynamic programming. Let odd_ways be the number of ways to get an odd sum and even_ways be the number of ways to get an even sum. We can build the solution iteratively. For a path of length 1, there's 1 way to get an odd sum (weight 1) and 1 way to get an even sum (weight 2). For a path of length i, we can extend the paths of length i-1. An odd sum can be formed by taking an even sum and adding 1, or an odd sum and adding 2. Thus, the new number of odd ways is the sum of the old odd and even ways. The same logic applies to forming an new even sum.

1class Solution {2    private long MOD = 1000000007;3 4    public int numberOfWays(int[][] edges) {5        int n = edges.length + 1;6        if (n <= 1) return 0;7 8        List<List<Integer>> adj = new ArrayList<>();9        for (int i = 0; i <= n; i++) {10            adj.add(new ArrayList<>());11        }12        for (int[] edge : edges) {13            adj.get(edge[0]).add(edge[1]);14            adj.get(edge[1]).add(edge[0]);15        }16 17        int L = findMaxDepth(n, adj);18        if (L == 0) return 0;19 20        // DP state for path of length 121        long evenWays = 1; // assign weight 222        long oddWays = 1;  // assign weight 123 24        // Iterate from path length 2 to L25        for (int i = 2; i <= L; i++) {26            long nextEven = (evenWays + oddWays) % MOD;27            long nextOdd = (evenWays + oddWays) % MOD;28            evenWays = nextEven;29            oddWays = nextOdd;30        }31 32        return (int) oddWays;33    }34 35    private int findMaxDepth(int n, List<List<Integer>> adj) {36        Queue<int[]> queue = new LinkedList<>();37        queue.offer(new int[]{1, 0}); // {node, depth}38        boolean[] visited = new boolean[n + 1];39        visited[1] = true;40        int maxDepth = 0;41 42        while (!queue.isEmpty()) {43            int[] current = queue.poll();44            int u = current[0];45            int depth = current[1];46            maxDepth = Math.max(maxDepth, depth);47 48            for (int v : adj.get(u)) {49                if (!visited[v]) {50                    visited[v] = true;51                    queue.offer(new int[]{v, depth + 1});52                }53            }54        }55        return maxDepth;56    }57}

Complexity

Time

O(n + L) = O(n). `O(n)` for finding the max depth `L`. The DP calculation involves a loop that runs `L-1` times, taking `O(L)` time. Since `L < n`, the total complexity is dominated by `O(n)`.

Space

O(n). `O(n)` for the adjacency list and BFS data structures. The DP calculation itself is optimized to use `O(1)` space.

Trade-offs

Pros

Efficient enough to pass the given constraints.
Systematically builds the solution from smaller subproblems.

Cons

The calculation part is linear in L, which is slightly less optimal than the final approach, though both are dominated by the O(n) graph traversal.

This is the most efficient and elegant approach. It leverages a combinatorial insight to find the answer directly without iteration or recursion on the path length. After finding the max depth L, the problem reduces to a single mathematical calculation which can be performed very quickly.

Algorithm

Build Graph & Find Max Depth (L): Use BFS to find the maximum depth L of the tree, which corresponds to the number of edges in the path.
Combinatorial Insight: Realize that the total path cost is odd if and only if the number of edges with weight 1 (an odd weight) is odd. The problem is thus equivalent to choosing an odd number of edges out of L to assign weight 1.
Mathematical Formula: The number of ways to choose an odd number of items from a set of L is given by the sum C(L, 1) + C(L, 3) + C(L, 5) + .... This sum is a known combinatorial identity that equals 2^(L-1) for L > 0.
Calculation: The number of ways is 2^(L-1). We need to compute this value modulo 10^9 + 7.
Modular Exponentiation: Use the efficient 'exponentiation by squaring' algorithm to calculate pow(2, L-1) under the given modulus. This computes the result in logarithmic time with respect to the exponent.

Walkthrough

The most direct solution comes from a mathematical observation. A sum of integers is odd if and only if it contains an odd number of odd integers. In our case, the weights are 1 (odd) and 2 (even). Therefore, for the total path cost to be odd, we must assign weight 1 to an odd number of edges on the path.

Let the path length be L. The total number of ways to assign weights to L edges is 2^L. For any path of length i, the number of ways to get an odd sum is equal to the number of ways to get an even sum, which is 2^(i-1). This can be proven by induction or by the DP relations from the previous approach. For a path of length L, the number of ways to get an odd sum is 2^(L-1).

We first find L using BFS, and then compute 2^(L-1) modulo 10^9 + 7 using modular exponentiation.

1class Solution {2    private long MOD = 1000000007;3 4    public int numberOfWays(int[][] edges) {5        int n = edges.length + 1;6        if (n <= 1) return 0;7 8        List<List<Integer>> adj = new ArrayList<>();9        for (int i = 0; i <= n; i++) {10            adj.add(new ArrayList<>());11        }12        for (int[] edge : edges) {13            adj.get(edge[0]).add(edge[1]);14            adj.get(edge[1]).add(edge[0]);15        }16 17        int L = findMaxDepth(n, adj);18        if (L == 0) return 0;19 20        return (int) power(2, L - 1);21    }22 23    private long power(long base, long exp) {24        long res = 1;25        base %= MOD;26        while (exp > 0) {27            if (exp % 2 == 1) res = (res * base) % MOD;28            base = (base * base) % MOD;29            exp /= 2;30        }31        return res;32    }33 34    private int findMaxDepth(int n, List<List<Integer>> adj) {35        Queue<int[]> queue = new LinkedList<>();36        queue.offer(new int[]{1, 0}); // {node, depth}37        boolean[] visited = new boolean[n + 1];38        visited[1] = true;39        int maxDepth = 0;40 41        while (!queue.isEmpty()) {42            int[] current = queue.poll();43            int u = current[0];44            int depth = current[1];45            maxDepth = Math.max(maxDepth, depth);46 47            for (int v : adj.get(u)) {48                if (!visited[v]) {49                    visited[v] = true;50                    queue.offer(new int[]{v, depth + 1});51                }52            }53        }54        return maxDepth;55    }56}

Complexity

Time

O(n + log L) = O(n). `O(n)` for finding the max depth `L` via BFS. The modular exponentiation takes `O(log L)` time. The overall complexity is dominated by the `O(n)` graph traversal.

Space

O(n). `O(n)` for the adjacency list and BFS data structures. The modular exponentiation function uses `O(1)` space for the iterative version.

Trade-offs

Pros

Most efficient method, especially the calculation part.
Code is concise and clean.

Cons

Requires recognizing the underlying mathematical pattern or combinatorial identity.

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem