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

• Java
• C++
• Python
• Go
• TypeScript

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

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

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

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

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

All Problems

All Solutions