根据题目描述和参考解法,以下是 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),用于存储倍增表。

 

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐