3558. Number of Ways to Assign Edge Weights I LeetCode Solution

Last updated: July 31, 2025

In this guide, you will get 3558. Number of Ways to Assign Edge Weights I LeetCode Solution with the best time and space complexity. The solution to Number of Ways to Assign Edge Weights I problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Problem Statement
Complexity Analysis
Number of Ways to Assign Edge Weights I solution in C++
Number of Ways to Assign Edge Weights I solution in Java
Number of Ways to Assign Edge Weights I solution in Python
Additional Resources

3558. Number of Ways to Assign Edge Weights I LeetCode Solution image

Problem Statement of Number of Ways to Assign Edge Weights I

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.

Complexity Analysis

Time Complexity: O(n)
Space Complexity: O(n)

3558. Number of Ways to Assign Edge Weights I LeetCode Solution in C++

class Solution {
 public:
  int assignEdgeWeights(vector<vector<int>>& edges) {
    const int n = edges.size() + 1;
    vector<vector<int>> graph(n + 1);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].push_back(v);
      graph[v].push_back(u);
    }

    queue<int> q{{1}};
    vector<bool> seen(n + 1);
    seen[1] = true;
    int step = 0;

    for (step = 0; !q.empty(); ++step)
      for (int sz = q.size(); sz > 0; --sz) {
        const int u = q.front();
        q.pop();
        for (const int v : graph[u])
          if (!seen[v]) {
            q.push(v);
            seen[v] = true;
          }
      }

    return step > 0 ? modPow(2, step - 2) : 0;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % kMod;
  }
};
/* code provided by PROGIEZ */

3558. Number of Ways to Assign Edge Weights I LeetCode Solution in Java

class Solution {
  public int assignEdgeWeights(int[][] edges) {
    final int n = edges.size() + 1;
    List<Integer>[] graph = new List[n + 1];
    Arrays.setAll(graph, i -> new ArrayList<>());

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(v);
      graph[v].add(u);
    }

    Queue<Integer> q = new ArrayDeque<>(List.of(1));
    boolean[] seen = new boolean[n + 1];
    seen[1] = true;

    int step = 0;
    for (step = 0; !q.isEmpty(); ++step)
      for (int sz = q.size(); sz > 0; --sz) {
        final int u = q.poll();
        for (final int v : graph[u])
          if (!seen[v]) {
            q.offer(v);
            seen[v] = true;
          }
      }

    return step > 0 ? modPow(2, step - 2) : 0;
  }

  private static final int MOD = 1_000_000_007;

  private int modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return (int) (x * modPow(x % MOD, (n - 1)) % MOD);
    return modPow(x * x % MOD, (n / 2)) % MOD;
  }
}
// code provided by PROGIEZ

3558. Number of Ways to Assign Edge Weights I LeetCode Solution in Python

class Solution:
  def assignEdgeWeights(self, edges: list[list[int]]) -> int:
    MOD = 1_000_000_007
    n = len(edges) + 1
    graph = [[] for _ in range(n + 1)]

    for u, v in edges:
      graph[u].append(v)
      graph[v].append(u)

    q = collections.deque([1])
    seen = {1}

    step = 0
    while q:
      for _ in range(len(q)):
        u = q.popleft()
        for v in graph[u]:
          if v not in seen:
            q.append(v)
            seen.add(v)
      step += 1

    return pow(2, step - 2, MOD) if step > 0 else 0
 # code by PROGIEZ

Additional Resources

Happy Coding! Keep following PROGIEZ for more updates and solutions.

Table of Contents