Efficient Solution for “Number of Ways to Assign Edge Weights I” on LeetCode


Problem Summary
You are given a tree with n
nodes rooted at node 1
, and each edge initially has a weight of 0
. The task is to assign each edge a weight of either 1
or 2
. The cost of a path is defined as the sum of edge weights along the path.
Use the main crux 2^(depth-1), because there are two ways to do it, either takes 1 or 2, and depth is a dependent variable here.
possible ways for even or odd=> 2^(depth)
for odd or even, odd=even=2^(depth-1) ,depth =maximum depth.
Pick any node
x
at the maximum depth in the tree (i.e., the farthest from the root).Determine how many different ways you can assign weights
1
or2
to only the edges on the path from the node1
to nodex
such that the total cost of the path is odd.
Key Observations
Since weights can only be
1
or2
Each edge has two assignment options.If the depth (distance from the root to the selected node
x
) isd
Then there are2^d
total weight combinations.Out of these, exactly half will result in an odd total cost because adding 1 or 2 changes the parity predictably.
Therefore, the number of valid assignments is:
2^(d - 1) % MOD
, whereMOD = 1e9 + 7
.
Solution Approach
We need to:
Build the adjacency list of the tree from the edge list.
Perform a BFS starting from the node
1
to find the maximum depth in the tree.Use modular exponentiation to compute
2^(depth - 1)
.
Code Explanation
Here’s the clean and efficient C++ implementation:
class Solution {
public:
int MOD = 1e9 + 7;
int power(int base, int exp, int mod) {
long long result = 1;
long long b = base;
while (exp > 0) {
if (exp % 2 == 1) result = (result * b) % mod;
b = (b * b) % mod;
exp /= 2;
}
return result;
}
int assignEdgeWeights(vector<vector<int>>& edges) {
int n = edges.size(); // number of edges = n - 1
vector<vector<int>> graph(n + 2);
for (auto& edge : edges) {
graph[edge[0]].push_back(edge[1]);
graph[edge[1]].push_back(edge[0]);
}
vector<int> visited(n + 2, 0);
int maxDepth = 0;
queue<pair<int, int>> q;
q.push({0, 1}); // {depth, node}
visited[1] = 1;
while (!q.empty()) {
int size = q.size();
while (size--) {
auto [depth, node] = q.front(); q.pop();
maxDepth = max(maxDepth, depth);
for (int neighbor : graph[node]) {
if (!visited[neighbor]) {
visited[neighbor] = 1;
q.push({depth + 1, neighbor});
}
}
}
}
return power(2, maxDepth - 1, MOD);
}
};
Time and Space Complexity
TC:
O(n)
wheren
is the number of nodes, due to BFS traversal.SC:
O(n)
for storing the graph and the visited array.
Conclusion
Simple based on permutations. two ways to do it 2^x ,x is dependent variable, which is here the maximum depth.
Subscribe to my newsletter
Read articles from drashy sesodia directly inside your inbox. Subscribe to the newsletter, and don't miss out.
Written by
