3558. Number of Ways to Assign Edge Weights I

Tree
BFS
DFS
Math

TIME

O(n)

SPACE

O(n)

Intuition

Only the parity of the path cost matters. A weight of 2 is even and never changes parity; a weight of 1 is odd and flips it. So a path's total cost is odd exactly when an odd number of its edges are assigned weight 1.

For a path with d edges, each edge is independently 1 or 2, giving 2^d assignments. Exactly half of them have an odd number of 1s, so the count of odd-cost assignments is 2^(d-1).

The node x is at maximum depth, so d is simply the maximum depth of the tree (the number of edges on the longest root-to-node path). The specific node chosen does not matter, only d does.

Approach

Build an adjacency list from edges.
BFS from node 1 to find the maximum depth d (edges from the root to the deepest node).
Return 2^(d-1) mod (10^9 + 7) using fast modular exponentiation.

Complexity

Time: O(n), one BFS over the tree
Space: O(n), adjacency list and BFS queue

Edge Cases

A single edge (n == 2) -> d = 1, answer is 2^0 = 1
A path-shaped tree -> d = n - 1, the largest possible depth
Large d -> use modular exponentiation to stay within the modulus

Code

from collections import deque
from typing import List


class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        """
        3558. Number of Ways to Assign Edge Weights I
        Time: O(n)
        Space: O(n)
        """
        MOD = 10**9 + 7
        n = len(edges) + 1

        adj = [[] for _ in range(n + 1)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)

        # BFS from the root (node 1) to find the maximum depth in edges.
        max_depth = 0
        visited = [False] * (n + 1)
        visited[1] = True
        queue = deque([(1, 0)])

        while queue:
            node, depth = queue.popleft()
            max_depth = max(max_depth, depth)
            for neighbor in adj[node]:
                if not visited[neighbor]:
                    visited[neighbor] = True
                    queue.append((neighbor, depth + 1))

        return pow(2, max_depth - 1, MOD)

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;

/**
 * 3558. Number of Ways to Assign Edge Weights I
 * Time: O(n)
 * Space: O(n)
 */
class Solution {
    private static final int MOD = 1_000_000_007;

    public int assignEdgeWeights(int[][] edges) {
        int n = edges.length + 1;

        List<List<Integer>> adj = new ArrayList<>();
        for (int i = 0; i <= n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }

        // BFS from the root (node 1) to find the maximum depth in edges.
        int maxDepth = 0;
        boolean[] visited = new boolean[n + 1];
        visited[1] = true;
        Deque<int[]> queue = new ArrayDeque<>();
        queue.offer(new int[] {1, 0});

        while (!queue.isEmpty()) {
            int[] current = queue.poll();
            int node = current[0];
            int depth = current[1];
            maxDepth = Math.max(maxDepth, depth);
            for (int neighbor : adj.get(node)) {
                if (!visited[neighbor]) {
                    visited[neighbor] = true;
                    queue.offer(new int[] {neighbor, depth + 1});
                }
            }
        }

        return (int) modPow(2, maxDepth - 1, MOD);
    }

    private long modPow(long base, long exp, long mod) {
        long result = 1;
        base %= mod;
        while (exp > 0) {
            if ((exp & 1) == 1) {
                result = result * base % mod;
            }
            base = base * base % mod;
            exp >>= 1;
        }
        return result;
    }
}