3559. Number of Ways to Assign Edge Weights II

  • Tree
  • Binary Lifting
  • Depth-First Search
  • Math
TIME
O((n + q) log n)
SPACE
O(n log n)
problem on leetcode ↗source on github ↗

Intuition

As in the easy version, only parity matters: weight 2 keeps parity, weight 1 flips it, so a path's cost is odd exactly when an odd number of its edges get weight 1. For a path with d edges there are 2^d assignments and exactly half are odd, giving 2^(d-1) (and 0 when d = 0, i.e. u == v).

So every query reduces to: how many edges lie on the path between u and v? In a tree that distance is

d = depth[u] + depth[v] - 2 * depth[lca(u, v)]

The only real work is answering many LCA queries quickly, which binary lifting handles in O(log n) each after O(n log n) preprocessing.

Approach

  1. Build an adjacency list and BFS from the root (node 1) to record depth[v] and the immediate parent up[0][v].
  2. Fill the binary-lifting table: up[k][v] = up[k-1][up[k-1][v]] for each level k.
  3. Precompute pow2[i] = 2^i mod (10^9 + 7) up to n.
  4. For each query [u, v]:
    • find lca(u, v) via binary lifting,
    • compute d = depth[u] + depth[v] - 2 * depth[lca],
    • append pow2[d - 1] if d >= 1, else 0.

Complexity

  • Time: O((n + q) log n), table build plus one LCA per query
  • Space: O(n log n), the binary-lifting ancestor table

Edge Cases

  • Query with u == v -> d = 0, answer is 0
  • Query where one node is an ancestor of the other -> LCA is that ancestor
  • Deepest possible path (d = n - 1) -> still O(log n) per query, modular power handles the size
  • Repeated/duplicate queries -> each answered independently in O(log n)

Code

from collections import deque
from typing import List


class Solution:
    def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        """
        3559. Number of Ways to Assign Edge Weights II
        Time: O((n + q) log n)
        Space: O(n log 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)

        LOG = max(1, n.bit_length())
        up = [[0] * (n + 1) for _ in range(LOG)]
        depth = [0] * (n + 1)

        # BFS from root to set depth and immediate parent (up[0]).
        visited = [False] * (n + 1)
        visited[1] = True
        queue = deque([1])
        while queue:
            node = queue.popleft()
            for neighbor in adj[node]:
                if not visited[neighbor]:
                    visited[neighbor] = True
                    depth[neighbor] = depth[node] + 1
                    up[0][neighbor] = node
                    queue.append(neighbor)

        # Binary lifting table.
        for k in range(1, LOG):
            up_k, up_prev = up[k], up[k - 1]
            for v in range(1, n + 1):
                up_k[v] = up_prev[up_prev[v]]

        pow2 = [1] * (n + 1)
        for i in range(1, n + 1):
            pow2[i] = pow2[i - 1] * 2 % MOD

        def lca(a: int, b: int) -> int:
            if depth[a] < depth[b]:
                a, b = b, a
            diff = depth[a] - depth[b]
            for k in range(LOG):
                if (diff >> k) & 1:
                    a = up[k][a]
            if a == b:
                return a
            for k in range(LOG - 1, -1, -1):
                if up[k][a] != up[k][b]:
                    a = up[k][a]
                    b = up[k][b]
            return up[0][a]

        answer = []
        for u, v in queries:
            d = depth[u] + depth[v] - 2 * depth[lca(u, v)]
            answer.append(pow2[d - 1] if d >= 1 else 0)

        return answer