Number of Ways to Assign Edge Weights II

Hard

#3169Time: `O(Q * N)`, where `Q` is the number of queries and `N` is the number of nodes. For each query, a BFS traversal might visit all `N` nodes and `N-1` edges in the worst case.Space: `O(N)` for the adjacency list and the data structures used in BFS (queue and visited array).

Patterns

Math Dynamic Programming

Algorithms

Depth-First Search

Data structures

Array 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.

You are given a 2D integer array queries. For each queries[i] = [u_i, v_i], determine the number of ways to assign weights to edges in the path such that the cost of the path between u_i and v_i is odd.

Return an array answer, where answer[i] is the number of valid assignments for queries[i].

Since the answer may be large, apply modulo 10⁹ + 7 to each answer[i].

Note: For each query, disregard all edges not in the path between node u_i and v_i.

Example 1:

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

Output: [0,1]

Explanation:

Query [1,1]: The path from Node 1 to itself consists of no edges, so the cost is 0. Thus, the number of valid assignments is 0.
Query [1,2]: 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]], queries = [[1,4],[3,4],[2,5]]

Output: [2,1,4]

Explanation:

Query [1,4]: 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.
Query [3,4]: The path from Node 3 to Node 4 consists of one edge (3 → 4). Assigning weight 1 makes the cost odd, while 2 makes it even. Thus, the number of valid assignments is 1.
Query [2,5]: The path from Node 2 to Node 5 consists of three edges (2 → 1, 1 → 3, and 3 → 5). Assigning (1,2,2), (2,1,2), (2,2,1), or (1,1,1) makes the cost odd. Thus, the number of valid assignments is 4.

Constraints:

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

Approaches

2 approaches with complexity analysis and trade-offs.

For each query, this approach finds the path between the two nodes by performing a Breadth-First Search (BFS) starting from one node until the other is found. The length of this path is then used to calculate the number of ways.

Algorithm

Construct an adjacency list for the tree from the edges input.
Initialize an answer array for the queries.
For each query (u, v):
- a. If u == v, the answer is 0. Continue to the next query.
- b. Perform a BFS starting from u to find the shortest path length to v.
  - i. Create a queue and add (u, 0).
  - ii. Create a visited set/array and add u.
  - iii. While the queue is not empty, dequeue (currentNode, distance).
  - iv. If currentNode == v, the path length k is distance. Break the BFS.
  - v. For each neighbor of currentNode, if not visited, mark as visited and enqueue (neighbor, distance + 1).
- c. Calculate power(2, k - 1, 10^9 + 7) using modular exponentiation.
- d. Store the result in the answer array.
Return the answer array.

Walkthrough

The core idea is to treat each query independently. First, the edges array is used to build an adjacency list representation of the tree. For each query (u, v), a BFS is initiated from node u. A queue stores pairs of (node, distance_from_u), and a visited array prevents redundant exploration. When node v is reached, its corresponding distance from u gives the path length, let's call it k.

The problem asks for the number of ways to assign weights {1, 2} to the k edges in the path such that the total cost is odd. The total cost is odd if and only if an odd number of edges have a weight of 1. For a path of length k, the number of ways to choose an odd number of edges to have weight 1 is given by the sum of binomial coefficients C(k, 1) + C(k, 3) + ..., which is equal to 2^(k-1) for k >= 1. If k = 0 (i.e., u == v), the path has no edges, the cost is 0 (even), so there are 0 ways.

The final answer for a path of length k is 2^(k-1) modulo 10^9 + 7.

Here is the Java implementation:

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

Complexity

Time

`O(Q * N)`, where `Q` is the number of queries and `N` is the number of nodes. For each query, a BFS traversal might visit all `N` nodes and `N-1` edges in the worst case.

Space

`O(N)` for the adjacency list and the data structures used in BFS (queue and visited array).

Trade-offs

Pros

Simple to understand and implement.
Does not require complex data structures or precomputation.

Cons

Highly inefficient for large inputs.
The O(Q * N) complexity will lead to a 'Time Limit Exceeded' error given the problem constraints.

This approach significantly speeds up the process by precomputing information about the tree. It uses the formula dist(u, v) = depth(u) + depth(v) - 2 * depth(lca(u, v)) to find the path length. The depths of all nodes and a data structure for finding the Lowest Common Ancestor (LCA) efficiently are prepared in a preprocessing step.

Algorithm

Preprocessing:
- a. Define MAX_LOG_N based on n (e.g., ceil(log2(n))).
- b. Build the adjacency list from edges.
- c. Initialize depth and up arrays. up[i][j] will store the 2^j-th ancestor of i.
- d. Perform a DFS from the root (node 1) to populate depth and up[i][0] (the parent) for all nodes.
- e. Fill the rest of the up table using the relation up[i][j] = up[up[i][j-1]][j-1].
- f. Precompute powers of 2 modulo 10^9 + 7 and store them in an array for O(1) lookup.
Query Processing:
- a. For each query (u, v):
  - i. If u == v, the answer is 0.
  - ii. Find lca_node = lca(u, v) using the binary lifting table.
  - iii. Calculate path length k = depth[u] + depth[v] - 2 * depth[lca_node].
  - iv. The number of ways is pow2[k-1].
  - v. Store the result.
- b. Return the array of results.
LCA Function lca(u, v):
- 1. Ensure u is the deeper node by swapping if necessary.
- 1. Lift u up by depth[u] - depth[v] levels so it's at the same depth as v, using binary jumps.
- 1. If u == v, then u is the LCA. Return u.
- 1. Lift u and v up together using binary jumps until their parents are the same.
- 1. The parent, up[u][0], is the LCA.

Walkthrough

The key insight is that the path length k between two nodes u and v in a tree can be calculated quickly if we know their depths and the depth of their LCA. The formula is k = depth(u) + depth[v] - 2 * depth(lca(u, v)). The number of ways to get an odd path cost is then 2^(k-1).

To use this formula, we first preprocess the tree:

Adjacency List: Build a graph representation.
DFS Traversal: Perform a single DFS from the root (node 1) to compute the depth of each node and the immediate parent of each node.
Binary Lifting Table: Use the parent information to build a table, up[i][j], which stores the 2^j-th ancestor of node i. This is constructed dynamically: up[i][0] is the parent of i, and up[i][j] = up[up[i][j-1]][j-1]. This table allows us to jump up the tree in powers of two, enabling fast LCA queries.

After this one-time O(N log N) preprocessing, each query (u, v) can be answered efficiently:

Find their LCA using the binary lifting table in O(log N) time.
Calculate the path length k using the depth formula in O(1).
The result is 2^(k-1) (or 0 if k=0), which can be found in O(1) if powers of 2 are precomputed.

Here is the Java implementation:

1class Solution {2    private List<List<Integer>> adj;3    private int[] depth;4    private int[][] up;5    private int MAX_LOG_N;6    private long[] pow2;7    private final int MOD = 1_000_000_007;8 9    public int[] numberOfWays(int n, int[][] edges, int[][] queries) {10        adj = new ArrayList<>();11        for (int i = 0; i <= n; i++) {12            adj.add(new ArrayList<>());13        }14        for (int[] edge : edges) {15            adj.get(edge[0]).add(edge[1]);16            adj.get(edge[1]).add(edge[0]);17        }18 19        MAX_LOG_N = (int) (Math.log(n) / Math.log(2)) + 1;20        depth = new int[n + 1];21        up = new int[n + 1][MAX_LOG_N];22 23        // Preprocessing24        dfs(1, 0, 0);25        for (int j = 1; j < MAX_LOG_N; j++) {26            for (int i = 1; i <= n; i++) {27                up[i][j] = up[up[i][j - 1]][j - 1];28            }29        }30        31        pow2 = new long[n];32        pow2[0] = 1;33        for (int i = 1; i < n; i++) {34            pow2[i] = (pow2[i - 1] * 2) % MOD;35        }36 37        // Process queries38        int[] result = new int[queries.length];39        for (int i = 0; i < queries.length; i++) {40            int u = queries[i][0];41            int v = queries[i][1];42            43            if (u == v) {44                result[i] = 0;45                continue;46            }47 48            int lcaNode = lca(u, v);49            int pathLength = depth[u] + depth[v] - 2 * depth[lcaNode];50            51            result[i] = (int) pow2[pathLength - 1];52        }53 54        return result;55    }56 57    private void dfs(int u, int p, int d) {58        depth[u] = d;59        up[u][0] = p;60        for (int v : adj.get(u)) {61            if (v != p) {62                dfs(v, u, d + 1);63            }64        }65    }66 67    private int lca(int u, int v) {68        if (depth[u] < depth[v]) {69            int temp = u;70            u = v;71            v = temp;72        }73 74        for (int j = MAX_LOG_N - 1; j >= 0; j--) {75            if (depth[u] - (1 << j) >= depth[v]) {76                u = up[u][j];77            }78        }79 80        if (u == v) {81            return u;82        }83 84        for (int j = MAX_LOG_N - 1; j >= 0; j--) {85            if (up[u][j] != up[v][j]) {86                u = up[u][j];87                v = up[v][j];88            }89        }90        return up[u][0];91    }92}

Complexity

Time

`O(N * log N + Q * log N)`. The preprocessing step takes `O(N * log N)` to build the binary lifting table. Each of the `Q` queries takes `O(log N)` to find the LCA.

Space

`O(N * log N)`. This is dominated by the space required for the `up` table used in binary lifting.

Trade-offs

Pros

Highly efficient for a large number of queries.
It's a standard and powerful technique for tree-based path problems.

Cons

More complex to implement than the brute-force approach.
Requires understanding of LCA and binary lifting.
The space complexity is higher due to the binary lifting table.

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