3559. Number of Ways to Assign Edge Weights II

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.

You are given a 2D integer array queries. For each queries[i] = [ui, vi], determine the number of ways to assign weights to edges in the path such that the cost of the path between ui and vi is odd.

Return an array answer, where answer[i] is the number of valid assignments for queries[i].

Since the answer may be large, apply modulo 10^9 + 7 to each answer[i].

Note: For each query, disregard all edges not in the path between node ui and vi.

Examples

Example 1

Input: edges = [[1,2]], queries = [[1,1],[1,2]]
Output: [0,1]

Explanation:

Query [1,1]: The path from Node 1 to itself consists of no edges, so the cost is 0. Thus, the number of valid assignments is 0.
Query [1,2]: 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]], queries = [[1,4],[3,4],[2,5]]
Output: [2,1,4]

Explanation:

Query [1,4]: 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.
Query [3,4]: The path from Node 3 to Node 4 consists of one edge (3 → 4). Assigning weight 1 makes the cost odd, while 2 makes it even. Thus, the number of valid assignments is 1.
Query [2,5]: The path from Node 2 to Node 5 consists of three edges (2 → 1, 1 → 3, and 3 → 5). Assigning (1,2,2), (2,1,2), (2,2,1), or (1,1,1) makes the cost odd. Thus, the number of valid assignments is 4.

Constraints

2 <= n <= 10^5
edges.length == n - 1
edges[i] == [ui, vi]
1 <= queries.length <= 10^5
queries[i] == [ui, vi]
1 <= ui, vi <= n
edges represents a valid tree.

Explanation

each edge on the path gets weight 1 or 2, and we want the count of assignments where the total cost is odd.

first observation — the actual values don’t matter at all, only the parity. weight 1 flips the parity of the sum, weight 2 keeps it. so the question is really: in how many ways can we pick odd/even for each edge so that the sum comes out odd?

Counting the ways

let’s say the path between u and v has d edges. take the first d - 1 edges and assign them whatever you like — that’s 2^(d-1) combinations. now look at the parity of what those d - 1 edges sum to:

if it’s already odd → the last edge must be 2 (even) to keep it odd
if it’s even → the last edge must be 1 (odd) to make it odd

either way, exactly one of the two choices works for the last edge. so every one of those 2^(d-1) prefixes extends to exactly one valid assignment:

ways(d) = 2^(d-1)

special case: if u == v the path has no edges, the cost is 0 which is even, so the answer is 0.

you can verify with example 2 — the path from 2 to 5 has 3 edges, and 2^2 = 4 matches the four assignments listed.

Getting the path length

so the whole problem reduces to: for up to 10^5 queries, find the number of edges on the path between u and v in a tree. that’s a classic LCA job:

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

with binary lifting we answer each lca query in O(log n) after O(n log n) preprocessing. I wrote a separate blog explaining LCA and binary lifting from scratch — link in the resources if you want the full picture.

for 2^(d-1) mod 1e9+7 use binary exponentiation. one small detail in my pow — it returns 0 when the exponent is negative, which conveniently handles the u == v case (d - 1 = -1) without writing an extra if.

overall complexity is O((n + q) log n).

Code

#define ll long long
#define all(x) x.begin(),x.end()
const int mod = 1e9+7;

struct BinaryLift {
    int n, lg;
    vector<int> depth;
    vector<vector<int>> adj, up;
    BinaryLift (int _n) : n (_n), lg (__lg (n - 1) + 1), depth (n), adj (n), up (lg, vector<int> (n, 0) ) {}

    void addEdge (int u, int v) {
        adj[u].push_back (v);
        adj[v].push_back (u);
    }

    void init (int r) {
        dfs (r);
        for (int i = 1; i < lg; i++)
            for (int j = 0; j < n; j++)
                if (up[i - 1][j] != -1)
                    up[i][j] = up[i - 1][up[i - 1][j]];
    }
    void dfs (int u) {
        for (int v : adj[u])
            if (v != up[0][u]) {
                depth[v] = depth[u] + 1;
                up[0][v] = u;
                dfs (v);
            }
    }

    int lca (int u, int v) {
        if (depth[u] < depth[v])
            swap (u, v);

        for (int i = lg - 1; i >= 0; i--)
            if (depth[u] - (1 << i) >= depth[v])
                u = up[i][u];

        if (u == v)
            return u;

        for (int i = lg - 1; i >= 0; i--)
            if (up[i][u] != up[i][v]) {
                u = up[i][u];
                v = up[i][v];
            }

        return up[0][u];
    }

    int dist (int u, int v) {
        return depth[u] + depth[v] - 2 * depth[lca (u, v)];
    }
};

int pow(int a,int b){
    if(b<0) return 0;
    int res = 1;
    for(int p = b; p; p>>=1, a = (a*1ll*a)%mod)
        if(p&1) res = (res*1ll*a)%mod;
    return res%mod;
}

class Solution {
public:
    vector<int> assignEdgeWeights(vector<vector<int>>& edges, vector<vector<int>>& queries) {
        int n = edges.size()+1;
        BinaryLift graph (n + 1);
        vector<int>ans;
        for(auto &v:edges){
            graph.addEdge(v[0],v[1]);
        }
        graph.init(1);
        for(auto &q:queries){
            auto d = graph.dist(q[0],q[1]);
            ans.push_back(pow(2,d-1));
        }
        return ans;
    }
};