Medium

题目描述

给你一个由 n 个顶点组成的无根加权树，顶点从 0 到 n - 1 编号，表示服务器。给你一个数组 edges，其中 edges[i] = [ai, bi, weighti] 表示顶点 ai 和 bi 之间有一条权重为 weighti 的双向边。同时给你一个整数 signalSpeed。

如果满足下述条件，则认为两个服务器 a 和 b 可以通过服务器 c 连通：

  • a < b，且 a != c，b != c
  • 从 c 到 a 的距离可以被 signalSpeed 整除
  • 从 c 到 b 的距离可以被 signalSpeed 整除
  • 从 c 到 a 的路径与从 c 到 b 的路径不共享任何边

返回长度为 n 的整数数组 count，其中 count[i] 是可以通过服务器 i 连通的服务器对数量。

示例 1：

输入：edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1
输出：[0,4,6,6,4,0]
解释：由于 signalSpeed 是 1，count[c] 等于从 c 开始且不共享边的路径对数量。
对于给定的路径图，count[c] 等于 c 左边的服务器数量乘以 c 右边的服务器数量。

示例 2：

输入：edges = [[0,6,3],[6,5,3],[0,3,1],[3,2,7],[3,1,6],[3,4,2]], signalSpeed = 3
输出：[2,0,0,0,0,0,2]
解释：通过服务器 0，有 2 对可连通的服务器：(4, 5) 和 (4, 6)。
通过服务器 6，有 2 对可连通的服务器：(4, 5) 和 (0, 5)。
可以证明，除了服务器 0 和 6 之外，没有两个服务器可以通过其他服务器连通。

提示：

  • 2 <= n <= 1000
  • edges.length == n - 1
  • edges[i].length == 3
  • 0 <= ai, bi < n
  • edges[i] = [ai, bi, weighti]
  • 1 <= weighti <= 10^6
  • 1 <= signalSpeed <= 10^6
  • 输入保证 edges 表示一个有效的树

解题思路

这道题的核心思路是对于每个可能的中转服务器，计算通过它能连通的服务器对数量。

算法思路：

  1. 建图：首先根据边信息构建邻接表表示的树结构。

  2. 以每个节点为根进行 DFS：对于每个节点作为中转服务器的情况，我们需要：

    • 以该节点为根，对其每个子树进行 DFS
    • 计算每个子树中距离根节点的距离能被 signalSpeed 整除的节点数量

  3. 计算可连通对数：关键观察是，两个服务器要通过中转服务器连通，它们必须：

    • 位于不同的子树中（保证路径不共享边）
    • 到中转服务器的距离都能被 signalSpeed 整除

  4. 组合计数：如果中转服务器有多个子树，分别包含 count[0], count[1], ..., count[k] 个满足条件的节点，那么可连通的对数就是从这些节点中任选两个且来自不同子树的方案数。

具体计算公式：对于子树 i，它与其他所有子树的节点组成的对数为 count[i] * (total - count[i])，其中 total 是所有子树中满足条件的节点总数。但这样会重复计算，所以最终结果需要除以 2。

更简洁的计算方法是：∑(count[i] * count[j]) 其中 i < j。

代码实现

class Solution {
public:
    vector<vector<pair<int, int>>> graph;
    
    int dfs(int node, int parent, int dist, int signalSpeed) {
        int count = (dist % signalSpeed == 0) ? 1 : 0;
        for (auto& edge : graph[node]) {
            int next = edge.first;
            int weight = edge.second;
            if (next != parent) {
                count += dfs(next, node, dist + weight, signalSpeed);
            }
        }
        return count;
    }
    
    vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
        int n = edges.size() + 1;
        graph.resize(n);
        
        // 建图
        for (auto& edge : edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            graph[u].push_back({v, w});
            graph[v].push_back({u, w});
        }
        
        vector<int> result(n, 0);
        
        // 以每个节点为中转服务器
        for (int root = 0; root < n; root++) {
            vector<int> subtreeCounts;
            
            // 对每个子树进行 DFS
            for (auto& edge : graph[root]) {
                int child = edge.first;
                int weight = edge.second;
                int count = dfs(child, root, weight, signalSpeed);
                subtreeCounts.push_back(count);
            }
            
            // 计算可连通的对数
            int total = 0;
            for (int count : subtreeCounts) {
                result[root] += count * total;
                total += count;
            }
        }
        
        return result;
    }
};

class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]:
        n = len(edges) + 1
        graph = [[] for _ in range(n)]
        
        # 建图
        for u, v, w in edges:
            graph[u].append((v, w))
            graph[v].append((u, w))
        
        def dfs(node, parent, dist):
            count = 1 if dist % signalSpeed == 0 else 0
            for next_node, weight in graph[node]:
                if next_node != parent:
                    count += dfs(next_node, node, dist + weight)
            return count
        
        result = [0] * n
        
        # 以每个节点为中转服务器
        for root in range(n):
            subtree_counts = []
            
            # 对每个子树进行 DFS
            for child, weight in graph[root]:
                count = dfs(child, root, weight)
                subtree_counts.append(count)
            
            # 计算可连通的对数
            total = 0
            for count in subtree_counts:
                result[root] += count * total
                total += count
        
        return result

public class Solution {
    private List<List<(int, int)>> graph;
    
    private int Dfs(int node, int parent, int dist, int signalSpeed) {
        int count = (dist % signalSpeed == 0) ? 1 : 0;
        foreach (var edge in graph[node]) {
            int next = edge.Item1;
            int weight = edge.Item2;
            if (next != parent) {
                count += Dfs(next, node, dist + weight, signalSpeed);
            }
        }
        return count;
    }
    
    public int[] CountPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.Length + 1;
        graph = new List<List<(int, int)>>(n);
        for (int i = 0; i < n; i++) {
            graph.Add(new List<(int, int)>());
        }
        
        // 建图
        foreach (var edge in edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            graph[u].Add((v, w));
            graph[v].Add((u, w));
        }
        
        int[] result = new int[n];
        
        // 以每个节点为中转服务器
        for (int root = 0; root < n; root++) {
            var subtreeCounts = new List<int>();
            
            // 对每个子树进行 DFS
            foreach (var edge in graph[root]) {
                int child = edge.Item1;
                int weight = edge.Item2;
                int count = Dfs(child, root, weight, signalSpeed);
                subtreeCounts.Add(count);
            }
            
            // 计算可连通的对数
            int total = 0;
            foreach (int count in subtreeCounts) {
                result[root] += count * total;
                total += count;
            }
        }
        
        return result;
    }
}

var countPairsOfConnectableServers = function(edges, signalSpeed) {
    const n = edges.length + 1;
    const graph = Array(n).fill().map(() => []);
    
    for (const [a, b, w] of edges) {
        graph[a].push([b, w]);
        graph[b].push([a, w]);
    }
    
    const result = Array(n).fill(0);
    
    for (let c = 0; c < n; c++) {
        const subtreeCounts = [];
        
        for (const [neighbor, weight] of graph[c]) {
            const count = dfs(neighbor, c, weight, graph, signalSpeed);
            subtreeCounts.push(count);
        }
        
        let pairs = 0;
        for (let i = 0; i < subtreeCounts.length; i++) {
            for (let j = i + 1; j < subtreeCounts.length; j++) {
                pairs += subtreeCounts[i] * subtreeCounts[j];
            }
        }
        result[c] = pairs;
    }
    
    return result;
};

function dfs(node, parent, distFromRoot, graph, signalSpeed) {
    let count = distFromRoot % signalSpeed === 0 ? 1 : 0;
    
    for (const [neighbor, weight] of graph[node]) {
        if (neighbor !== parent) {
            count += dfs(neighbor, node, distFromRoot + weight, graph, signalSpeed);
        }
    }
    
    return count;
}
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型大小
时间复杂度O(n²)
空间复杂度O(n)

时间复杂度分析：

  • 对于每个节点作为根节点，都需要进行一次完整的树遍历
  • 每次遍历的时间复杂度为 O(n)
  • 总共需要进行 n 次遍历，所以总时间复杂度为 O(n²)

空间复杂度分析：

  • 图的邻接表存储需要 O(n) 空间
  • DFS 递归栈的最大深度为树的高度，最坏情况下为 O(n)
  • 其他辅助数组也是 O(n) 空间
  • 总空间复杂度为 O(n)

相关题目