Number of ways to assign edge weights i

Math Tree Depth-First Search

53.2% acceptance

Feb 25, 2026

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.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn assign_edge_weights(edges: Vec<Vec<i32>>) -> i32 {
    // Path from root to deepest node has d edges (d = max depth).
    // Number of ways to assign weights 1 or 2 to d edges such that sum is odd = 2^(d-1).
    // For d=0: 0 ways. For d=1: 1 way. For d>=1: 2^(d-1).
    
    let n = edges.len() + 1;
    let mut adj: Vec<Vec<usize>> = vec![vec![]; n + 1]; // nodes 1..=n
    for e in &edges {
      let (u, v) = (e[0] as usize, e[1] as usize);
      adj[u].push(v);
      adj[v].push(u);
    }
    
    // BFS from node 1 to find max depth
    let mut depth = vec![0i32; n + 1];
    let mut visited = vec![false; n + 1];
    let mut queue = std::collections::VecDeque::new();
    queue.push_back(1usize);
    visited[1] = true;
    let mut max_depth = 0;
    while let Some(u) = queue.pop_front() {
      if depth[u] > max_depth { max_depth = depth[u]; }
      for &v in &adj[u] {
        if !visited[v] {
          visited[v] = true;
          depth[v] = depth[u] + 1;
          queue.push_back(v);
        }
      }
    }
    
    const MOD: i64 = 1_000_000_007;
    if max_depth == 0 { return 0; }
    // 2^(max_depth - 1) mod MOD
    let mut result = 1i64;
    let exp = max_depth as u64 - 1;
    let mut base = 2i64;
    let mut e = exp;
    while e > 0 {
      if e & 1 == 1 { result = result * base % MOD; }
      base = base * base % MOD;
      e >>= 1;
    }
    result as i32
  }
}