Kimi LeetCode 3559. 给边赋权值的方案数 II Go实现

根据题目描述和参考解法，以下是 LeetCode 3559. 给边赋权值的方案数 II 的 Go 实现。

解题思路

核心观察：每条边只能赋权值 1 或 2。
- 权值为 1 的边对路径代价贡献奇数（1）
- 权值为 2 的边对路径代价贡献偶数（2）

设路径上有 `d` 条边，要使总代价为奇数，需要选择奇数条边赋值为 1（其余赋值为 2）。

对于 `d` 条边，选奇数个赋值为 1 的方案数 = 2^{d-1}。

因此问题转化为：对每次查询，求两点之间路径上的边数 `d`，答案为 2^{d-1} \bmod (10^9+7)。当 u=v 时，答案为 0。

求路径边数：使用 LCA（最近公共祖先）算法，路径边数 d = depth[u] + depth[v] - 2 \times depth[lca]。

Go 实现

```go
package main

import (
"math"
)

func assignEdgeWeights(edges [][]int, queries [][]int) []int {
const MOD int = 1_000_000_007
n := len(edges) + 1
LOG := int(math.Ceil(math.Log2(float64(n)))) + 1

// 建图
graph := make([][]int, n+1)
for _, e := range edges {
u, v := e[0], e[1]
graph[u] = append(graph[u], v)
graph[v] = append(graph[v], u)
}

// depth[u]: 节点u到根节点1的深度（边数）
depth := make([]int, n+1)
// parent[k][u]: 节点u的2^k级祖先
parent := make([][]int, LOG)
for i := range parent {
parent[i] = make([]int, n+1)
for j := range parent[i] {
parent[i][j] = -1
}
}

// DFS预处理深度和一级祖先
var dfs func(u, p int)
dfs = func(u, p int) {
parent[0][u] = p
for _, v := range graph[u] {
if v != p {
depth[v] = depth[u] + 1
dfs(v, u)
}
}
}
dfs(1, -1)

// 倍增预处理各级祖先
for k := 1; k < LOG; k++ {
for v := 1; v <= n; v++ {
if parent[k-1][v] != -1 {
parent[k][v] = parent[k-1][parent[k-1][v]]
}
}
}

// LCA查询
lca := func(u, v int) int {
if depth[u] < depth[v] {
u, v = v, u
}
// 将u提升到与v同一深度
for k := LOG - 1; k >= 0; k-- {
if parent[k][u] != -1 && depth[parent[k][u]] >= depth[v] {
u = parent[k][u]
}
}
if u == v {
return u
}
// 同时向上提升u和v
for k := LOG - 1; k >= 0; k-- {
if parent[k][u] != -1 && parent[k][u] != parent[k][v] {
u = parent[k][u]
v = parent[k][v]
}
}
return parent[0][u]
}

// 快速幂
modPow := func(x, n int) int {
res := 1
x = x % MOD
for n > 0 {
if n&1 == 1 {
res = res * x % MOD
}
x = x * x % MOD
n >>= 1
}
return res
}

ans := make([]int, len(queries))
for i, q := range queries {
u, v := q[0], q[1]
if u == v {
ans[i] = 0
continue
}
a := lca(u, v)
d := depth[u] + depth[v] - 2*depth[a] // 路径边数
ans[i] = modPow(2, d-1)
}
return ans
}
```

复杂度分析

- 时间复杂度：O(n \log n + q \log n)，其中预处理 DFS 和倍增为 O(n \log n)，每次查询 LCA 为 O(\log n)。
- 空间复杂度：O(n \log n)，用于存储倍增表。