Hard

题目描述

给你一个有 n 个节点的无向加权树,节点编号从 0 到 n - 1。树用长度为 n - 1 的二维整数数组 edges 表示,其中 edges[i] = [ui, vi, wi] 表示节点 ui 和 vi 之间有一条权重为 wi 的边。

另外,给你一个二维整数数组 queries,其中 queries[j] = [src1j, src2j, destj]。

返回一个长度等于 queries.length 的数组 answer,其中 answer[j] 是子树的最小总权重,使得可以使用这个子树中的边从 src1j 和 src2j 都能到达 destj。

这里的子树是原树的任何连通的节点和边的子集,形成一个有效的树。

示例 1:

输入:edges = [[0,1,2],[1,2,3],[1,3,5],[1,4,4],[2,5,6]], queries = [[2,3,4],[0,2,5]]
输出:[12,11]
解释:
蓝色边代表产生最优答案的子树之一。
answer[0]:选择的子树的总权重确保从 src1 = 2 和 src2 = 3 到 dest = 4 的路径是 3 + 5 + 4 = 12。
answer[1]:选择的子树的总权重确保从 src1 = 0 和 src2 = 2 到 dest = 5 的路径是 2 + 3 + 6 = 11。

示例 2:

输入:edges = [[1,0,8],[0,2,7]], queries = [[0,1,2]]
输出:[15]
解释:
answer[0]:选择的子树的总权重确保从 src1 = 0 和 src2 = 1 到 dest = 2 的路径是 8 + 7 = 15。

约束条件:

  • 3 <= n <= 10^5
  • edges.length == n - 1
  • edges[i].length == 3
  • 0 <= ui, vi < n
  • 1 <= wi <= 10^4
  • 1 <= queries.length <= 10^5
  • queries[j].length == 3
  • 0 <= src1j, src2j, destj < n
  • src1j, src2j 和 destj 两两不同
  • 输入保证 edges 表示一个有效的树

解题思路

这道题的核心是找到连接三个节点的最小生成树。对于任意三个节点 a、b、c,连接它们的最小子树权重为 (d(a,b) + d(b,c) + d(c,a)) / 2,其中 d(x,y) 表示节点 x 和 y 之间的距离。

解题思路:

  1. 构建树结构:使用邻接表存储树的边和权重信息。

  2. 预处理距离计算

    • 使用 DFS 预计算每个节点到根节点的距离
    • 实现二进制提升(Binary Lifting)算法来快速计算最低公共祖先(LCA)
    • 利用公式 d(x,y) = dist[x] + dist[y] - 2 * dist[LCA(x,y)] 计算任意两点距离
  3. 处理查询

    • 对于每个查询 [src1, src2, dest],计算三点之间的两两距离
    • 应用公式 (d(src1,src2) + d(src1,dest) + d(src2,dest)) / 2 得到最小子树权重
  4. 优化要点

    • 二进制提升预处理时间复杂度 O(n log n)
    • 每次 LCA 查询时间复杂度 O(log n)
    • 总体时间复杂度 O((n + q) log n),其中 q 是查询数量

这种方法的关键在于理解树上任意三点的最小连通子图就是包含这三点的 Steiner 树,其权重恰好等于三点两两距离之和的一半。

代码实现

class Solution {
public:
    vector<vector<pair<int, int>>> adj;
    vector<long long> dist;
    vector<vector<int>> up;
    int LOG;
    
    void dfs(int u, int p, long long d) {
        dist[u] = d;
        up[u][0] = p;
        
        for (int i = 1; i < LOG; i++) {
            if (up[u][i-1] != -1) {
                up[u][i] = up[up[u][i-1]][i-1];
            }
        }
        
        for (auto& edge : adj[u]) {
            int v = edge.first;
            int w = edge.second;
            if (v != p) {
                dfs(v, u, d + w);
            }
        }
    }
    
    int lca(int u, int v) {
        if (dist[u] < dist[v]) swap(u, v);
        
        // Bring u to the same level as v
        for (int i = LOG - 1; i >= 0; i--) {
            if (up[u][i] != -1 && dist[up[u][i]] >= dist[v]) {
                u = up[u][i];
            }
        }
        
        if (u == v) return u;
        
        // Binary search for LCA
        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 long distance(int u, int v) {
        int l = lca(u, v);
        return dist[u] + dist[v] - 2 * dist[l];
    }
    
    vector<int> minimumWeight(vector<vector<int>>& edges, vector<vector<int>>& queries) {
        int n = edges.size() + 1;
        LOG = ceil(log2(n)) + 1;
        
        adj.resize(n);
        dist.resize(n);
        up.assign(n, vector<int>(LOG, -1));
        
        for (auto& edge : edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            adj[u].push_back({v, w});
            adj[v].push_back({u, w});
        }
        
        dfs(0, -1, 0);
        
        vector<int> result;
        for (auto& query : queries) {
            int src1 = query[0], src2 = query[1], dest = query[2];
            
            long long d1 = distance(src1, src2);
            long long d2 = distance(src1, dest);
            long long d3 = distance(src2, dest);
            
            result.push_back((d1 + d2 + d3) / 2);
        }
        
        return result;
    }
};
class Solution:
    def minimumWeight(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        n = len(edges) + 1
        LOG = max(1, (n - 1).bit_length())
        
        adj = [[] for _ in range(n)]
        for u, v, w in edges:
            adj[u].append((v, w))
            adj[v].append((u, w))
        
        dist = [0] * n
        up = [[-1] * LOG for _ in range(n)]
        
        def dfs(u, p, d):
            dist[u] = d
            up[u][0] = p
            
            for i in range(1, LOG):
                if up[u][i-1] != -1:
                    up[u][i] = up[up[u][i-1]][i-1]
            
            for v, w in adj[u]:
                if v != p:
                    dfs(v, u, d + w)
        
        def lca(u, v):
            if dist[u] < dist[v]:
                u, v = v, u
            
            # Bring u to the same level as v
            for i in range(LOG-1, -1, -1):
                if up[u][i] != -1 and dist[up[u][i]] >= dist[v]:
                    u = up[u][i]
            
            if u == v:
                return u
            
            # Binary search for LCA
            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]
        
        def distance(u, v):
            l = lca(u, v)
            return dist[u] + dist[v] - 2 * dist[l]
        
        dfs(0, -1, 0)
        
        result = []
        for src1, src2, dest in queries:
            d1 = distance(src1, src2)
            d2 = distance(src1, dest)
            d3 = distance(src2, dest)
            result.append((d1 + d2 + d3) // 2)
        
        return result
public class Solution {
    private List<List<(int, int)>> adj;
    private long[] dist;
    private int[,] up;
    private int LOG;
    
    private void Dfs(int u, int p, long d) {
        dist[u] = d;
        up[u, 0] = p;
        
        for (int i = 1; i < LOG; i++) {
            if (up[u, i-1] != -1) {
                up[u, i] = up[up[u, i-1], i-1];
            } else {
                up[u, i] = -1;
            }
        }
        
        foreach (var edge in adj[u]) {
            int v = edge.Item1;
            int w = edge.Item2;
            if (v != p) {
                Dfs(v, u, d + w);
            }
        }
    }
    
    private int Lca(int u, int v) {
        if (dist[u] < dist[v]) {
            (u, v) = (v, u);
        }
        
        // Bring u to the same level as v
        for (int i = LOG - 1; i >= 0; i--) {
            if (up[u, i] != -1 && dist[up[u, i]] >= dist[v]) {
                u = up[u, i];
            }
        }
        
        if (u == v) return u;
        
        // Binary search for LCA
        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];
    }
    
    private long Distance(int u, int v) {
        int l = Lca(u, v);
        return dist[u] + dist[v] - 2 * dist[l];
    }
    
    public int[] MinimumWeight(int[][] edges, int[][] queries) {
        int n = edges.Length + 1;
        LOG = (int)Math.Ceiling(Math.Log2(n)) + 1;
        
        adj = new List<List<(int, int)>>();
        for (int i = 0; i < n; i++) {
            adj.Add(new List<(int, int)>());
        }
        
        dist = new long[n];
        up = new int[n, LOG];
        
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < LOG; j++) {
                up[i, j] = -1;
            }
        }
        
        foreach (var edge in edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            adj[u].Add((v, w));
            adj[v].Add((u, w));
        }
        
        Dfs(0, -1, 0);
        
        int[] result = new int[queries.Length];
        for (int i = 0; i < queries.Length; i++) {
            int src1 = queries[i][0], src2 = queries[i][1], dest = queries[i][2];
            
            long d1 = Distance(src1, src2);
            long d2 = Distance(src1, dest);
            long d3 = Distance(src2, dest);
            
            result[i] = (int)((d1 + d2 + d3) / 2);
        }
        
        return result;
    }
}
var minimumWeight = function(edges, queries) {
    const n = edges.length + 1;
    const graph = Array(n).fill().map(() => []);
    
    for (const [u, v, w] of edges) {
        graph[u].push([v, w]);
        graph[v].push([u, w]);
    }
    
    const getPath = (start, end) => {
        if (start === end) return [];
        
        const parent = Array(n).fill(-1);
        const queue = [start];
        parent[start] = start;
        
        while (queue.length > 0) {
            const node = queue.shift();
            
            for (const [neighbor, weight] of graph[node]) {
                if (parent[neighbor] === -1) {
                    parent[neighbor] = node;
                    queue.push(neighbor);
                    
                    if (neighbor === end) {
                        const path = [];
                        let curr = end;
                        while (curr !== start) {
                            const prev = parent[curr];
                            for (const [next, w] of graph[prev]) {
                                if (next === curr) {
                                    path.push([prev, curr, w]);
                                    break;
                                }
                            }
                            curr = prev;
                        }
                        return path.reverse();
                    }
                }
            }
        }
        return [];
    };
    
    const result = [];
    
    for (const [src1, src2, dest] of queries) {
        const path1 = getPath(src1, dest);
        const path2 = getPath(src2, dest);
        
        const edgeSet = new Set();
        
        for (const [u, v, w] of path1) {
            const key = u < v ? `${u}-${v}` : `${v}-${u}`;
            edgeSet.add(key);
        }
        
        for (const [u, v, w] of path2) {
            const key = u < v ? `${u}-${v}` : `${v}-${u}`;
            edgeSet.add(key);
        }
        
        let totalWeight = 0;
        const edgeWeights = new Map();
        
        for (const [u, v, w] of edges) {
            const key = u < v ? `${u}-${v}` : `${v}-${u}`;
            edgeWeights.set(key, w);
        }
        
        for (const edgeKey of edgeSet) {
            totalWeight += edgeWeights.get(edgeKey);
        }
        
        result.push(totalWeight);
    }
    
    return result;
};

复杂度分析

指标复杂度
时间-
空间-