3559. 边权重分配的方案数II

Hard

题目描述

有一个包含n个节点的无向树，节点标号从1到n，根节点为1。树用长度为n-1的二维整数数组edges表示，其中edges[i] = [ui, vi]表示节点ui和vi之间有一条边。

初始时，所有边的权重都为0。你必须给每条边分配权重1或2。

任意两个节点u和v之间路径的代价是连接它们路径上所有边的权重总和。

给定一个二维整数数组queries。对于每个queries[i] = [ui, vi]，确定有多少种给边分配权重的方法，使得节点ui和vi之间路径的代价为奇数。

返回一个数组answer，其中answer[i]是queries[i]的有效分配方案数。

由于答案可能很大，对每个answer[i]取模109 + 7。

注意：对于每个查询，忽略不在节点ui和vi之间路径上的所有边。

示例1:

输入: edges = [[1,2]], queries = [[1,1],[1,2]]
输出: [0,1]
解释:
- 查询[1,1]: 从节点1到自己的路径不包含任何边，代价为0。因此有效分配数为0。
- 查询[1,2]: 从节点1到节点2的路径包含一条边(1→2)。分配权重1使代价为奇数，分配权重2使代价为偶数。因此有效分配数为1。

示例2:

输入: edges = [[1,2],[1,3],[3,4],[3,5]], queries = [[1,4],[3,4],[2,5]]
输出: [2,1,4]

约束条件:

2 <= n <= 10^5
edges.length == n - 1
edges[i] == [ui, vi]
1 <= queries.length <= 10^5
queries[i] == [ui, vi]
1 <= ui, vi <= n
edges表示一个有效的树

解题思路

这是一个树上路径问题，核心思路是动态规划 + 最近公共祖先(LCA)。

思路分析：

路径代价的奇偶性：要使路径代价为奇数，路径上权重为1的边数必须为奇数。每条边可以分配权重1或2，这本质上是一个组合问题。
动态规划状态：设dp[len][parity]表示长度为len的路径中，有parity个奇数权重边(权重为1)的方案数。状态转移为：
- dp[i][0] = dp[i-1][1] + dp[i-1][0]（当前边选权重2或1）
- dp[i][1] = dp[i-1][0] + dp[i-1][1]（当前边选权重1或2）
路径长度计算：使用LCA算法快速计算任意两点间的距离。对于查询[u,v]，路径长度为depth[u] + depth[v] - 2*depth[lca(u,v)]。
优化观察：通过数学分析可以发现，对于长度为k的路径，使代价为奇数的方案数为2^(k-1)（当k>0时），k=0时方案数为0。

算法步骤：

构建树的邻接表表示
使用DFS预处理每个节点的深度和父节点信息，构建LCA查询结构
对每个查询，计算路径长度并应用公式2^(k-1)

这种方法避免了复杂的DP计算，直接利用数学性质求解，时间复杂度更优。

代码实现

class Solution {
public:
    const int MOD = 1e9 + 7;
    
    vector<int> assignEdgeWeights(vector<vector<int>>& edges, vector<vector<int>>& queries) {
        int n = edges.size() + 1;
        vector<vector<int>> adj(n + 1);
        
        for (auto& edge : edges) {
            adj[edge[0]].push_back(edge[1]);
            adj[edge[1]].push_back(edge[0]);
        }
        
        // LCA preprocessing
        int LOG = 20;
        vector<vector<int>> up(n + 1, vector<int>(LOG));
        vector<int> depth(n + 1);
        
        function<void(int, int)> dfs = [&](int u, int p) {
            up[u][0] = p;
            for (int i = 1; i < LOG; i++) {
                up[u][i] = up[up[u][i-1]][i-1];
            }
            for (int v : adj[u]) {
                if (v != p) {
                    depth[v] = depth[u] + 1;
                    dfs(v, u);
                }
            }
        };
        
        depth[1] = 0;
        dfs(1, 1);
        
        auto lca = [&](int u, int v) {
            if (depth[u] < depth[v]) swap(u, v);
            
            int diff = depth[u] - depth[v];
            for (int i = 0; i < LOG; i++) {
                if ((diff >> i) & 1) {
                    u = up[u][i];
                }
            }
            
            if (u == v) return u;
            
            for (int i = LOG - 1; i >= 0; i--) {
                if (up[u][i] != up[v][i]) {
                    u = up[u][i];
                    v = up[v][i];
                }
            }
            
            return up[u][0];
        };
        
        auto power = [&](long long base, int exp) {
            long long result = 1;
            while (exp > 0) {
                if (exp & 1) result = (result * base) % MOD;
                base = (base * base) % MOD;
                exp >>= 1;
            }
            return (int)result;
        };
        
        vector<int> result;
        for (auto& query : queries) {
            int u = query[0], v = query[1];
            int pathLength = depth[u] + depth[v] - 2 * depth[lca(u, v)];
            
            if (pathLength == 0) {
                result.push_back(0);
            } else {
                result.push_back(power(2, pathLength - 1));
            }
        }
        
        return result;
    }
};

class Solution:
    def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        MOD = 10**9 + 7
        n = len(edges) + 1
        adj = [[] for _ in range(n + 1)]
        
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        # LCA preprocessing
        LOG = 20
        up = [[0] * LOG for _ in range(n + 1)]
        depth = [0] * (n + 1)
        
        def dfs(u, p):
            up[u][0] = p
            for i in range(1, LOG):
                up[u][i] = up[up[u][i-1]][i-1]
            
            for v in adj[u]:
                if v != p:
                    depth[v] = depth[u] + 1
                    dfs(v, u)
        
        dfs(1, 1)
        
        def lca(u, v):
            if depth[u] < depth[v]:
                u, v = v, u
            
            diff = depth[u] - depth[v]
            for i in range(LOG):
                if (diff >> i) & 1:
                    u = up[u][i]
            
            if u == v:
                return u
            
            for i in range(LOG - 1, -1, -1):
                if up[u][i] != up[v][i]:
                    u = up[u][i]
                    v = up[v][i]
            
            return up[u][0]
        
        result = []
        for u, v in queries:
            path_length = depth[u] + depth[v] - 2 * depth[lca(u, v)]
            
            if path_length == 0:
                result.append(0)
            else:
                result.append(pow(2, path_length - 1, MOD))
        
        return result

public class Solution {
    private const int MOD = 1000000007;
    
    public int[] AssignEdgeWeights(int[][] edges, int[][] queries) {
        int n = edges.Length + 1;
        var adj = new List<int>[n + 1];
        for (int i = 0; i <= n; i++) {
            adj[i] = new List<int>();
        }
        
        foreach (var edge in edges) {
            adj[edge[0]].Add(edge[1]);
            adj[edge[1]].Add(edge[0]);
        }
        
        // LCA preprocessing
        int LOG = 20;
        var up = new int[n + 1][];
        for (int i = 0; i <= n; i++) {
            up[i] = new int[LOG];
        }
        var depth = new int[n + 1];
        
        void Dfs(int u, int p) {
            up[u][0] = p;
            for (int i = 1; i < LOG; i++) {
                up[u][i] = up[up[u][i-1]][i-1];
            }
            
            foreach (int v in adj[u]) {
                if (v != p) {
                    depth[v] = depth[u] + 1;
                    Dfs(v, u);
                }
            }
        }
        
        Dfs(1, 1);
        
        int Lca(int u, int v) {
            if (depth[u] < depth[v]) {
                (u, v) = (v, u);
            }
            
            int diff = depth[u] - depth[v];
            for (int i = 0; i < LOG; i++) {
                if ((diff >> i & 1) == 1) {
                    u = up[u][i];
                }
            }
            
            if (u == v) return u;
            
            for (int i = LOG - 1; i >= 0; i--) {
                if (up[u][i] != up[v][i]) {
                    u = up[u][i];
                    v = up[v][i];
                }
            }
            
            return up[u][0];
        }
        
        long Power(long baseNum, int exp) {
            long result = 1;
            while (exp > 0) {
                if ((exp & 1) == 1) {
                    result = (result * baseNum) % MOD;
                }
                baseNum = (baseNum * baseNum) % MOD;
                exp >>= 1;
            }
            return result;
        }
        
        var result = new int[queries.Length];
        for (int i = 0; i < queries.Length; i++) {
            int u = queries[i][0], v = queries[i][1];
            int pathLength = depth[u] + depth[v] - 2 * depth[Lca(u, v)];
            
            if (pathLength == 0) {
                result[i] = 0;
            } else {
                result[i] = (int)Power(2, pathLength - 1);
            }
        }
        
        return result;
    }
}

var assignEdgeWeights = function(edges, queries) {
    const MOD = 1000000007;
    const n = edges.length + 1;
    
    // Build adjacency list
    const adj = Array(n + 1).fill().map(() => []);
    for (const [u, v] of edges) {
        adj[u].push(v);
        adj[v].push(u);
    }
    
    // Build parent relationships using DFS from root 1
    const parent = Array(n + 1).fill(-1);
    const visited = Array(n + 1).fill(false);
    
    function dfs(node) {
        visited[node] = true;
        for (const neighbor of adj[node]) {
            if (!visited[neighbor]) {
                parent[neighbor] = node;
                dfs(neighbor);
            }
        }
    }
    
    dfs(1);
    
    // Find path between two nodes
    function findPath(u, v) {
        const pathU = [];
        const pathV = [];
        
        let curr = u;
        while (curr !== -1) {
            pathU.push(curr);
            curr = parent[curr];
        }
        
        curr = v;
        while (curr !== -1) {
            pathV.push(curr);
            curr = parent[curr];
        }
        
        // Find LCA
        const setU = new Set(pathU);
        let lca = -1;
        for (const node of pathV) {
            if (setU.has(node)) {
                lca = node;
                break;
            }
        }
        
        // Build path from u to v through lca
        const path = [];
        curr = u;
        while (curr !== lca) {
            path.push(curr);
            curr = parent[curr];
        }
        path.push(lca);
        
        const pathToV = [];
        curr = v;
        while (curr !== lca) {
            pathToV.push(curr);
            curr = parent[curr];
        }
        
        for (let i = pathToV.length - 1; i >= 0; i--) {
            path.push(pathToV[i]);
        }
        
        return path;
    }
    
    const result = [];
    
    for (const [u, v] of queries) {
        if (u === v) {
            result.push(0);
            continue;
        }
        
        const path = findPath(u, v);
        const numEdges = path.length - 1;
        
        if (numEdges === 0) {
            result.push(0);
            continue;
        }
        
        // For k edges, we need odd sum
        // Each edge can be 1 or 2
        // We need odd number of odd weights (weight 1)
        // Number of ways = sum of C(k, i) for odd i from 1 to k
        // This equals 2^(k-1) when k > 0
        
        let ways = 1;
        for (let i = 0; i < numEdges - 1; i++) {
            ways = (ways * 2) % MOD;
        }
        
        result.push(ways);
    }
    
    return result;
};

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

项目	复杂度
时间复杂度	O(n log n + q log n)
空间复杂度	O(n log n)

其中 n 是节点数，q 是查询数。时间复杂度包括 O(n log n) 的LCA预处理和 O(q log n) 的查询处理。空间复杂度主要由LCA的二进制提升表占用。

题目描述#

解题思路#

代码实现#

复杂度分析#

题目描述

解题思路

代码实现

复杂度分析