Welcome to Subscribe On Youtube

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

Solution 2: Explicit Stack + Mathematics

Thinking

The path from $1$ to a deepest leaf has $d$ edges. Each edge is $1$ or $2$, and the cost is odd iff an odd number of them are $1$. Other edges are free. The number of odd-sized subsets of $d$ elements is $2^{d-1}$ ($0$ when $d = 0$).

With $n \le 10^5$, recursing from the root into each child reaches depth $n$ on a chain and overflows the call stack. That depth is only the number of edges from the root, so sibling order does not change the maximum.

An explicit stack of $(node, parent, depth)$ therefore starts at node $1$ with depth $0$. Popping a frame updates the maximum depth, and every neighbor other than the parent is pushed with the depth increased by one. A fast power then evaluates $2^{d-1}$.

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

An explicit stack then computes the depth $d$ of the tree. Each frame is $(node, parent, depth)$, starting from node $1$ at depth $0$. Popping a frame updates $d$ with the current depth, and every neighbor other than the parent is pushed with the depth increased by one. The answer is the number of ways to choose an odd number of elements from $d$. According to a well-known combinatorial identity, that number 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.

  • 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;
    }
}


// Solution 2
class Solution {
    public int assignEdgeWeights(int[][] edges) {
        int n = edges.length + 1;
        List<Integer>[] 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);
        }
        int d = 0;
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {1, 0, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], dep = cur[2];
            d = Math.max(d, dep);
            for (int j : g[i]) {
                if (j != fa) {
                    stk.push(new int[] {j, i, dep + 1});
                }
            }
        }
        return (int) pow(2, d - 1, 1_000_000_007);
    }

    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;
    }
}


  • 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;
    }
};


// Solution 2
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);
        }
        int d = 0;
        vector<array<int, 3>> stk{{1, 0, 0}};
        while (!stk.empty()) {
            auto [i, fa, dep] = stk.back();
            stk.pop_back();
            d = max(d, dep);
            for (int j : g[i]) {
                if (j != fa) {
                    stk.push_back({j, i, dep + 1});
                }
            }
        }
        return pow(2, d - 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;
    }
};


  • 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)



# Solution 2
class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        n = len(edges) + 1
        g = [[] for _ in range(n + 1)]
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)
        stk = [(1, 0, 0)]
        d = 0
        while stk:
            i, fa, dep = stk.pop()
            d = max(d, dep)
            for j in g[i]:
                if j != fa:
                    stk.append((j, i, dep + 1))
        return pow(2, d - 1, 10**9 + 7)


  • 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
}


// Solution 2
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)
	}

	d := 0
	stk := [][3]int{{1, 0, 0}}
	for len(stk) > 0 {
		cur := stk[len(stk)-1]
		stk = stk[:len(stk)-1]
		i, fa, dep := cur[0], cur[1], cur[2]
		d = max(d, dep)
		for _, j := range g[i] {
			if j != fa {
				stk = append(stk, [3]int{j, i, dep + 1})
			}
		}
	}
	return pow(2, d-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
}


  • 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);
}



// Solution 2
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 stk: [number, number, number][] = [[1, 0, 0]];
    let d = 0;
    while (stk.length) {
        const [i, fa, dep] = stk.pop()!;
        d = Math.max(d, dep);
        for (const j of g[i]) {
            if (j !== fa) {
                stk.push([j, i, dep + 1]);
            }
        }
    }

    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, d - 1, mod);
}

All Problems

All Solutions