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),用于存储倍增表。