DeepSeek    LeetCode 3559. 给边赋权值的方案数 II Java实现

根据题目要求,对于树中任意两个节点 u 和 v,路径长度为 d(边数),使路径总代价为奇数的赋值方案数为 2^(d-1)。因此核心问题转化为快速求树上两点距离。

以下是 Java 实现,使用二进制提升(Binary Lifting) 求 LCA,预处理 O(n log n),单次查询 O(log n):

```java
class Solution {
private static final int MOD = 1_000_000_007;
private List<List<Integer>> graph;
private int[] depth;
private int[][] up;

public int[] assignEdgeWeights(int[][] edges, int[][] queries) {
int n = edges.length + 1;
// 1. 建图
graph = new ArrayList<>();
for (int i = 0; i < n; i++) graph.add(new ArrayList<>());
for (int[] e : edges) {
int u = e[0] - 1, v = e[1] - 1;
graph.get(u).add(v);
graph.get(v).add(u);
}

// 2. 预处理:深度 + 倍增表
int LOG = 32 - Integer.numberOfLeadingZeros(n); // ceil(log2(n)) + 1
depth = new int[n];
up = new int[n][LOG];
dfs(0, -1);

// 3. 预处理 2 的幂次
int[] pow2 = new int[n + 1];
pow2[0] = 1;
for (int i = 1; i <= n; i++) {
pow2[i] = (int)((long)pow2[i - 1] * 2 % MOD);
}

// 4. 处理每个查询
int q = queries.length;
int[] ans = new int[q];
for (int i = 0; i < q; i++) {
int u = queries[i][0] - 1, v = queries[i][1] - 1;
int lca = findLca(u, v);
int dist = depth[u] + depth[v] - 2 * depth[lca];
ans[i] = dist == 0 ? 0 : pow2[dist - 1];
}
return ans;
}

// DFS 构建深度和倍增表
private void dfs(int node, int parent) {
up[node][0] = parent == -1 ? node : parent;
for (int j = 1; j < up[0].length; j++) {
up[node][j] = up[up[node][j - 1]][j - 1];
}
for (int nei : graph.get(node)) {
if (nei == parent) continue;
depth[nei] = depth[node] + 1;
dfs(nei, node);
}
}

// 将节点向上提升 diff 步
private int lift(int node, int diff) {
for (int j = 0; diff > 0; j++, diff >>= 1) {
if ((diff & 1) == 1) node = up[node][j];
}
return node;
}

// 求 LCA
private int findLca(int u, int v) {
if (depth[u] > depth[v]) {
int tmp = u; u = v; v = tmp;
}
v = lift(v, depth[v] - depth[u]);
if (u == v) return u;

for (int j = up[0].length - 1; j >= 0; j--) {
if (up[u][j] != up[v][j]) {
u = up[u][j];
v = up[v][j];
}
}
return up[u][0];
}
}
```

复杂度分析

· 时间复杂度:预处理 O(n log n),每个查询 O(log n)
· 空间复杂度:O(n log n)

核心思路补充

为什么答案是 2^(d-1)?路径上 d 条边每条可选 1 或 2。要让总和为奇数,等价于从 d 条边中选奇数条赋值为 1(其余为 2)。d 个元素中选奇数个的组合数之和 = 2^(d-1)。