Skip to main content
Talentd

Number of Ways to Assign Edge Weights I - Solution & Explanation

MediumMathTreeDepth-First Search10 min readAsked at: Amazon, Meta
Practice this problem

Problem Statement

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.

Approach Overview

Problem Overview: You are given a tree where edge weights are not fixed. The task is to count how many valid ways you can assign weights to the edges while satisfying constraints on distances along tree paths. Because the structure is a tree, every pair of nodes has exactly one simple path, which allows the constraints to be analyzed using a single traversal.

Approach 1: Brute Force Enumeration (Exponential Time)

The most direct idea is to try every possible weight assignment for every edge and verify whether the resulting tree satisfies the required distance conditions. If each edge can take two possible weights, the total number of assignments is 2^(n-1) for a tree with n nodes. For each assignment, compute path distances and validate the constraints. This approach runs in O(2^n * n) time with O(n) space. It only works for very small trees and mainly helps build intuition about which edges actually influence the constraints.

Approach 2: Tree DFS + Combinatorics (O(n))

The key observation is that constraints on distances only affect edges that lie on specific paths in the tree. Because a tree has a unique path between any two nodes, you can analyze the structure using a single Depth-First Search. During DFS, track how many constrained nodes appear in each subtree. When a subtree contributes to a constrained path, the connecting edge becomes restricted; otherwise it remains free to choose from the allowed weight options.

Each unrestricted edge multiplies the number of valid assignments. If an edge has two valid weight choices, it contributes a factor of 2 to the total count. By counting how many edges remain unrestricted during the DFS traversal, the final result becomes a simple power calculation such as 2^k where k is the number of flexible edges. The traversal processes every node and edge once, giving O(n) time complexity and O(n) space for recursion and adjacency storage.

This method works because trees eliminate cycles, so each edge’s contribution can be determined independently once the subtree structure is known. DFS naturally exposes parent–child relationships and subtree counts, which makes it ideal for these kinds of constraint propagation problems in tree structures combined with mathematical counting.

Recommended for interviews: The DFS + combinatorics approach is the expected solution. Interviewers typically want to see that you recognize the unique-path property of trees, propagate constraints with a single traversal, and convert the remaining freedom into a mathematical counting formula. Mentioning the brute-force enumeration first shows you understand the search space, but deriving the O(n) DFS counting solution demonstrates stronger algorithmic reasoning.

Solution

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

Next, we use a function 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.

Code

PythonJavaC++GoTypeScript

Python

Java

C++

Go

TypeScript

Try this approach in the editor →

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
Brute Force Weight EnumerationO(2^n * n)O(n)Only for very small trees or for understanding the constraint behavior
DFS + Combinatorial CountingO(n)O(n)Optimal approach for large trees where constraints affect only specific paths

Video Solution

Number of Ways to Assign Edge Weights I | Simplified | Leetcode 3558 | codestorywithMIK • codestorywithMIK • 9,307 views views

Watch 9 more video solutions →

Frequently Asked Questions

Is Number of Ways to Assign Edge Weights I easy or hard?
The problem is typically classified as Medium difficulty. The tree traversal itself is straightforward, but the challenge is recognizing how path constraints restrict certain edges while leaving others independent. Once that observation is made, the counting step becomes simple.
Number of Ways to Assign Edge Weights I Python/Java solution
The implementation builds an adjacency list for the tree and runs a recursive or iterative DFS. While traversing, it tracks subtree constraints and counts edges that remain free to choose weights. The final answer is computed using modular exponentiation such as pow(2, k). The same logic works in Python, Java, C++, and Go.
How to solve Number of Ways to Assign Edge Weights I in O(n)?
Build the adjacency list of the tree and run a DFS from any root. During traversal, track how constraints propagate through subtrees and determine whether an edge is forced or still flexible. Count the number of edges that remain unrestricted and compute the number of assignments using a power function like 2^k. Because each edge is processed once, the algorithm runs in linear time.
What is the best approach for Number of Ways to Assign Edge Weights I?
The best approach uses a Depth-First Search (DFS) on the tree combined with combinatorial counting. DFS identifies which edges are constrained by path requirements and which edges remain flexible. Each flexible edge contributes independent choices, so the final answer becomes a power calculation such as 2^k. This produces an O(n) time solution.
Is Number of Ways to Assign Edge Weights I asked at Google/Amazon/Meta?
Tree counting problems that combine DFS with combinatorics are common in interviews at companies like Google, Amazon, and Meta. Variants often ask you to analyze constraints along paths in a tree and count valid configurations. Recognizing the unique-path property of trees is the main skill tested.
What data structure is used in Number of Ways to Assign Edge Weights I?
The main data structure is an adjacency list representing the tree. A DFS traversal explores the graph while maintaining subtree information. The algorithm also relies on simple mathematical counting to compute the total number of valid assignments.
What is the time complexity of Number of Ways to Assign Edge Weights I?
The optimal solution runs in O(n) time where n is the number of nodes in the tree. A single DFS traversal processes every node and edge once, while the final answer is computed using a simple exponentiation step. Space complexity is O(n) due to the adjacency list and recursion stack.

Ready to solve this problem?

Practice Number of Ways to Assign Edge Weights I with our built-in code editor and test cases.

Practice on FleetCode
Previous

Add Two Numbers

Next

Reverse Integer