Number of Ways to Assign Edge Weights II
Problem
An undirected tree has n nodes labelled 1..n, rooted at node 1, given by edges. Every edge must be assigned a weight of 1 or 2. The cost of a path is the sum of its edge weights. For each query [u, v], count the number of weight assignments to the edges on the path from u to v that make the path cost odd, modulo 109+7.
A weight 2 is even and a weight 1 is odd, so only the weight-1 edges flip the parity of the cost. The cost is odd exactly when an odd number of the m path edges get weight 1. Choosing an odd-sized subset of m items can be done in 2m−1 ways (and 0 ways when m = 0). So each answer is 2dist(u,v)−1, where dist is the number of edges between u and v.
edges = [[1,2],[1,3],[3,4],[3,5]], queries = [[1,4],[3,4],[2,5]][2, 1, 4]def assignEdgeWeights(edges, queries):
MOD = 10**9 + 7
n = len(edges) + 1
g = [[] for _ in range(n + 1)]
for u, v in edges: # build adjacency list
g[u].append(v)
g[v].append(u)
LOG = max(1, n.bit_length())
depth = [0] * (n + 1)
up = [[0] * (n + 1) for _ in range(LOG)] # up[k][v] = 2^k-th ancestor
seen = [False] * (n + 1)
seen[1] = True
stack = [1] # iterative DFS from root 1
while stack:
node = stack.pop()
for nb in g[node]:
if not seen[nb]:
seen[nb] = True
depth[nb] = depth[node] + 1
up[0][nb] = node # immediate parent
stack.append(nb)
for k in range(1, LOG): # binary-lifting table
for v in range(1, n + 1):
up[k][v] = up[k - 1][up[k - 1][v]]
def lca(u, v):
if depth[u] < depth[v]:
u, v = v, u
diff = depth[u] - depth[v] # lift u up to v's depth
for k in range(LOG):
if (diff >> k) & 1:
u = up[k][u]
if u == v:
return u
for k in range(LOG - 1, -1, -1): # rise together until parents meet
if up[k][u] != up[k][v]:
u, v = up[k][u], up[k][v]
return up[0][u]
ans = []
for u, v in queries:
w = lca(u, v)
dist = depth[u] + depth[v] - 2 * depth[w]
ans.append(pow(2, dist - 1, MOD) if dist >= 1 else 0)
return ansfunction assignEdgeWeights(edges, queries) {
const MOD = 1000000007n;
const n = edges.length + 1;
const g = Array.from({ length: n + 1 }, () => []);
for (const [u, v] of edges) { // build adjacency list
g[u].push(v);
g[v].push(u);
}
const LOG = Math.max(1, 32 - Math.clz32(n));
const depth = new Array(n + 1).fill(0);
const up = Array.from({ length: LOG }, () => new Array(n + 1).fill(0));
const seen = new Array(n + 1).fill(false);
seen[1] = true;
const stack = [1]; // iterative DFS from root 1
while (stack.length) {
const node = stack.pop();
for (const nb of g[node]) {
if (!seen[nb]) {
seen[nb] = true;
depth[nb] = depth[node] + 1;
up[0][nb] = node; // immediate parent
stack.push(nb);
}
}
}
for (let k = 1; k < LOG; k++) // binary-lifting table
for (let v = 1; v <= n; v++)
up[k][v] = up[k - 1][up[k - 1][v]];
const lca = (u, v) => {
if (depth[u] < depth[v]) [u, v] = [v, u];
let diff = depth[u] - depth[v]; // lift u up to v's depth
for (let k = 0; k < LOG; k++)
if ((diff >> k) & 1) u = up[k][u];
if (u === v) return u;
for (let k = LOG - 1; k >= 0; k--) // rise together
if (up[k][u] !== up[k][v]) { u = up[k][u]; v = up[k][v]; }
return up[0][u];
};
const power = (e) => { // 2^e mod p
let r = 1n, b = 2n;
while (e > 0) { if (e & 1) r = r * b % MOD; b = b * b % MOD; e >>= 1; }
return r;
};
const ans = [];
for (const [u, v] of queries) {
const w = lca(u, v);
const dist = depth[u] + depth[v] - 2 * depth[w];
ans.push(dist >= 1 ? Number(power(dist - 1)) : 0);
}
return ans;
}int[] assignEdgeWeights(int[][] edges, int[][] queries) {
final long MOD = 1_000_000_007L;
int n = edges.length + 1;
List<List<Integer>> g = new ArrayList<>();
for (int i = 0; i <= n; i++) g.add(new ArrayList<>());
for (int[] e : edges) { // build adjacency list
g.get(e[0]).add(e[1]);
g.get(e[1]).add(e[0]);
}
int LOG = Math.max(1, 32 - Integer.numberOfLeadingZeros(n));
int[] depth = new int[n + 1];
int[][] up = new int[LOG][n + 1]; // up[k][v] = 2^k-th ancestor
boolean[] seen = new boolean[n + 1];
seen[1] = true;
Deque<Integer> stack = new ArrayDeque<>();
stack.push(1); // iterative DFS from root 1
while (!stack.isEmpty()) {
int node = stack.pop();
for (int nb : g.get(node)) if (!seen[nb]) {
seen[nb] = true;
depth[nb] = depth[node] + 1;
up[0][nb] = node; // immediate parent
stack.push(nb);
}
}
for (int k = 1; k < LOG; k++) // binary-lifting table
for (int v = 1; v <= n; v++)
up[k][v] = up[k - 1][up[k - 1][v]];
int[] ans = new int[queries.length];
for (int i = 0; i < queries.length; i++) {
int u = queries[i][0], v = queries[i][1];
int w = lca(u, v, depth, up, LOG);
int dist = depth[u] + depth[v] - 2 * depth[w];
ans[i] = dist >= 1 ? (int) power(2, dist - 1, MOD) : 0;
}
return ans;
}
int lca(int u, int v, int[] depth, int[][] up, int LOG) {
if (depth[u] < depth[v]) { int t = u; u = v; v = t; }
int diff = depth[u] - depth[v]; // lift u up to v's depth
for (int k = 0; k < LOG; k++)
if (((diff >> k) & 1) == 1) u = up[k][u];
if (u == v) return u;
for (int k = LOG - 1; k >= 0; k--) // rise together
if (up[k][u] != up[k][v]) { u = up[k][u]; v = up[k][v]; }
return up[0][u];
}
long power(long b, long e, long MOD) { // b^e mod p
long r = 1;
while (e > 0) { if ((e & 1) == 1) r = r * b % MOD; b = b * b % MOD; e >>= 1; }
return r;
}vector<int> assignEdgeWeights(vector<vector<int>>& edges,
vector<vector<int>>& queries) {
const long long MOD = 1000000007LL;
int n = edges.size() + 1;
vector<vector<int>> g(n + 1);
for (auto& e : edges) { // build adjacency list
g[e[0]].push_back(e[1]);
g[e[1]].push_back(e[0]);
}
int LOG = max(1, 32 - __builtin_clz(n));
vector<int> depth(n + 1, 0);
vector<vector<int>> up(LOG, vector<int>(n + 1, 0)); // 2^k-th ancestor
vector<bool> seen(n + 1, false);
seen[1] = true;
vector<int> stk = {1}; // iterative DFS from root 1
while (!stk.empty()) {
int node = stk.back(); stk.pop_back();
for (int nb : g[node]) if (!seen[nb]) {
seen[nb] = true;
depth[nb] = depth[node] + 1;
up[0][nb] = node; // immediate parent
stk.push_back(nb);
}
}
for (int k = 1; k < LOG; k++) // binary-lifting table
for (int v = 1; v <= n; v++)
up[k][v] = up[k - 1][up[k - 1][v]];
auto lca = [&](int u, int v) {
if (depth[u] < depth[v]) swap(u, v);
int diff = depth[u] - depth[v]; // lift u up to v's depth
for (int k = 0; k < LOG; k++)
if ((diff >> k) & 1) u = up[k][u];
if (u == v) return u;
for (int k = LOG - 1; k >= 0; k--) // rise together
if (up[k][u] != up[k][v]) { u = up[k][u]; v = up[k][v]; }
return up[0][u];
};
auto power = [&](long long e) { // 2^e mod p
long long r = 1, b = 2;
while (e > 0) { if (e & 1) r = r * b % MOD; b = b * b % MOD; e >>= 1; }
return r;
};
vector<int> ans;
for (auto& q : queries) {
int w = lca(q[0], q[1]);
int dist = depth[q[0]] + depth[q[1]] - 2 * depth[w];
ans.push_back(dist >= 1 ? (int) power(dist - 1) : 0);
}
return ans;
}Explanation
The whole problem reduces to one observation about parity. Every edge gets weight 1 (odd) or weight 2 (even), so only the weight-1 edges change the parity of the path cost. A path of m edges has an odd cost exactly when an odd number of those edges are assigned weight 1.
The number of ways to pick an odd-sized subset out of m items is C(m,1) + C(m,3) + … = 2m−1 for m ≥ 1, and there is no way to make a 0-edge path odd, so the answer for that case is 0. Therefore answer = 2dist(u,v) − 1 mod (109+7) where dist is the number of edges on the path.
All that remains is computing dist(u, v) quickly. We root the tree at node 1 and run a single traversal to record each node's depth and immediate parent. Then dist(u, v) = depth[u] + depth[v] − 2·depth[lca(u, v)], where lca is the lowest common ancestor.
To answer many queries fast we precompute a binary-lifting table up[k][v] = the 2k-th ancestor of v. To find the LCA we first lift the deeper node up to the shallower node's depth, then raise both nodes together in powers-of-two jumps until their parents coincide. Each query then costs O(log n).
Powers of two are taken with fast modular exponentiation. Example: for the path 2 → 1 → 3 → 5 the distance is 3, so the count of odd-cost weightings is 22 = 4.