Skip to content

3558. Number of Ways to Assign Edge Weights I

Description

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] = [ui, vi] indicates that there is an edge between nodes ui and vi.

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 109 + 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 <= 105
  • edges.length == n - 1
  • edges[i] == [ui, vi]
  • 1 <= ui, vi <= n
  • edges represents a valid tree.

Solutions

Solution 1: DFS + Mathematics

First, we build an adjacency list \(g\) from the edges, where \(g[u]\) contains all neighbors of node \(u\).

Next, we use a function \(\textit{dfs}\) to compute the depth \(d\) of the tree. The answer is the number of ways to choose an odd number of elements from \(d\). According to a well-known combinatorial identity, the number of ways to choose an odd number of elements from a set of size \(d\) is \(2^{d-1}\). Therefore, we can compute the answer using fast exponentiation.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the number of nodes in the tree.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        def dfs(i: int, fa: int = 0) -> int:
            res = 0
            for j in g[i]:
                if j != fa:
                    res = max(res, dfs(j, i) + 1)
            return res

        n = len(edges) + 1
        g = [[] for _ in range(n + 1)]
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)
        d = dfs(1)
        return pow(2, d - 1, 10**9 + 7)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class Solution {
    private List<Integer>[] g;

    public int assignEdgeWeights(int[][] edges) {
        int n = edges.length + 1;
        g = new List[n + 1];
        Arrays.setAll(g, k -> new ArrayList<>());

        for (var e : edges) {
            int u = e[0];
            int v = e[1];
            g[u].add(v);
            g[v].add(u);
        }

        return (int) pow(2, dfs(1, 0) - 1, 1_000_000_007);
    }

    private int dfs(int i, int fa) {
        int res = 0;
        for (int j : g[i]) {
            if (j != fa) {
                res = Math.max(res, dfs(j, i) + 1);
            }
        }
        return res;
    }

    private long pow(long a, int n, int mod) {
        long res = 1;
        while (n > 0) {
            if ((n & 1) != 0) {
                res = res * a % mod;
            }
            a = a * a % mod;
            n >>= 1;
        }
        return res;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Solution {
public:
    int assignEdgeWeights(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<vector<int>> g(n + 1);

        for (auto& e : edges) {
            int u = e[0];
            int v = e[1];
            g[u].push_back(v);
            g[v].push_back(u);
        }

        auto dfs = [&](this auto&& dfs, int i, int fa) -> int {
            int res = 0;
            for (int j : g[i]) {
                if (j != fa) {
                    res = max(res, dfs(j, i) + 1);
                }
            }
            return res;
        };

        return pow(2, dfs(1, 0) - 1, 1000000007);
    }

private:
    long long pow(long long a, int n, int mod) {
        long long res = 1;
        while (n > 0) {
            if (n & 1) {
                res = res * a % mod;
            }
            a = a * a % mod;
            n >>= 1;
        }
        return res;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
func assignEdgeWeights(edges [][]int) int {
    const mod = 1_000_000_007

    n := len(edges) + 1
    g := make([][]int, n+1)

    for _, e := range edges {
        u, v := e[0], e[1]
        g[u] = append(g[u], v)
        g[v] = append(g[v], u)
    }

    var dfs func(int, int) int
    dfs = func(i, fa int) int {
        res := 0
        for _, j := range g[i] {
            if j != fa {
                res = max(res, dfs(j, i)+1)
            }
        }
        return res
    }

    return pow(2, dfs(1, 0)-1, mod)
}

func pow(a, n, mod int) int {
    res := 1
    for n > 0 {
        if n&1 > 0 {
            res = res * a % mod
        }
        a = a * a % mod
        n >>= 1
    }
    return res
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
function assignEdgeWeights(edges: number[][]): number {
    const mod = 1_000_000_007;
    const n = edges.length + 1;
    const g: number[][] = Array.from({ length: n + 1 }, () => []);

    for (const [u, v] of edges) {
        g[u].push(v);
        g[v].push(u);
    }

    const dfs = (i: number, fa: number): number => {
        let res = 0;
        for (const j of g[i]) {
            if (j !== fa) {
                res = Math.max(res, dfs(j, i) + 1);
            }
        }
        return res;
    };

    const pow = (a: number, n: number, mod: number): number => {
        let res = 1n;
        let x = BigInt(a);
        const m = BigInt(mod);

        while (n > 0) {
            if (n & 1) {
                res = (res * x) % m;
            }
            x = (x * x) % m;
            n >>= 1;
        }

        return Number(res);
    };

    return pow(2, dfs(1, 0) - 1, mod);
}

Comments