3558. Number of Ways to Assign Edge Weights I

Problem Statement

3558. Number of Ways to Assign Edge Weights I

Medium

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.

Java

Source file

class Solution {
    public int assignEdgeWeights(int[][] edges) {
        // Construct graph (adjacency list)
        int n = edges.length;
        List<Integer>[] adj = new List[n + 2];
        for (int i = 0; i < n + 2; i++)
            adj[i] = new ArrayList<>();
        for (int[] e : edges) {
            adj[e[0]].add(e[1]);
            adj[e[1]].add(e[0]);
        }
        // Find max depth via dfs
        int maxDepth = dfs(1, adj, 0);
        // Calc. odd parity of 1 weighted edges, rest 2, viz - nC1+nC3+nC5+...nCn_or_n-1
        // Above series equates to 2^(n-1) via binary expansion
        // Now instead of pow(2, n-1)%10e7, we will use fast exponentiation
        return fastExponentiation(2, maxDepth - 1);
    }

    int dfs(int node, List<Integer>[] adj, int parent) {
        int maxDepth = 0;
        for (int child : adj[node]) {
            if (child == parent)
                continue;
            maxDepth = Math.max(maxDepth, dfs(child, adj, node) + 1);
        }
        return maxDepth;
    }

    static final int MOD = 1_000_000_007;

    int fastExponentiation(int x, int y) { // calc x^n % 10e7
        long ans = 1;
        long base = x;
        while (y > 0) {
            if ((y & 1) == 1)
                ans = (ans * base) % MOD;
            base = (base * base) % MOD;
            y >>= 1;
        }
        return (int) ans;
    }
}