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

drashy sesodiadrashy sesodia
3 min read

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 or 2 to only the edges on the path from the node 1 to node x such that the total cost of the path is odd.


Key Observations

  1. Since weights can only be 1 or 2Each edge has two assignment options.

  2. If the depth (distance from the root to the selected node x) is dThen there are 2^d total weight combinations.

  3. Out of these, exactly half will result in an odd total cost because adding 1 or 2 changes the parity predictably.

  4. Therefore, the number of valid assignments is:
    2^(d - 1) % MOD, where MOD = 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) where n 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.

0
Subscribe to my newsletter

Read articles from drashy sesodia directly inside your inbox. Subscribe to the newsletter, and don't miss out.

Written by

drashy sesodia
drashy sesodia