3067. 在带权树网络中统计可连接服务器对数目

给你一棵无根带权树，树中总共有 n nn 个节点，分别表示 n nn 个服务器，服务器从 0 00 到 n−1 n - 1n−1 编号。同时给你一个数组 edges edgesedges ，其中 edges[i]=[ai,bi,weigh ti]edges[i] = [a_i, b_i, weight_i]edges[i]=[ai​,bi​,weighti​] 表示节点 ai a_iai​ 和 bi b_ibi​ 之间有一条双向边，边的权值为 weigh ti weight_iweighti​ 。再给你一个整数 signalSpeed signalSpeedsignalSpeed 。

如果两个服务器 a aa，b bb 和 c cc 满足以下条件，那么我们称服务器 a aa 和 b bb 是通过服务器 c cc 可连接的：

a<b a < ba<b ，a≠c a\ne ca=c 且 b≠c b\ne cb=c 。

从 c cc 到 a aa 的距离是可以被 signalSpeed signalSpeedsignalSpeed 整除的。

从 c cc 到 b bb 的距离是可以被 signalSpeed signalSpeedsignalSpeed 整除的。

从 c cc 到 b bb 的路径与从 c cc 到 a aa 的路径没有任何公共边。

请你返回一个长度为 n nn 的整数数组 count countcount ，其中 count[i]count[i]count[i] 表示通过服务器 i ii 可连接 的服务器对的 数目 。

示例 1：

Copy
输入：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 左边服务器数目乘以右边服务器数目。

示例 2：

Copy
输入：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 才可连接，所以其他服务器对应的可连接服务器对数目都为 0 。

提示：

2<=n<=1000 2 <= n <= 10002<=n<=1000

edges.length==n−1 edges.length == n - 1edges.length==n−1

edges[i].length==3 edges[i].length == 3edges[i].length==3

0<=ai,bi<n 0 <= a_i, b_i < n0<=ai​,bi​<n

edges[i]=[ai,bi,weigh ti]edges[i] = [a_i, b_i, weight_i]edges[i]=[ai​,bi​,weighti​]

1<=weigh ti<=1 06 1 <= weight_i <= 10^61<=weighti​<=106

1<=signalSpeed<=1 06 1 <= signalSpeed <= 10^61<=signalSpeed<=106

输入保证 edges edgesedges 构成一棵合法的树。

根枚举

a<b a < ba<b ，a≠c a \ne ca=c 且 b≠c b \ne cb=c 。

从 c cc 到 a aa 的距离是可以被 signalSpeed signalSpeedsignalSpeed 整除的。

从 c cc 到 b bb 的距离是可以被 signalSpeed signalSpeedsignalSpeed 整除的。

从 c cc 到 b bb 的路径与从 c cc 到 a aa 的路径没有任何公共边。

从可连接的条件中我们可以得到以下信息：

a 和 b 位置可以互换，只需要满足存在差值即可。（也可以理解为防止重复）

c 是将这个无向无环图分为两个部分的节点，因为只有这样 c 到 a 的路径与 c 到 b 的路径才没有任何公共边。

分析以下案例：

遍历到子树 1：

遍历到子树 2：

遍历到子树 3：

将每一次遍历的满足条件的结果相加即为当前遍历的根节点的可连接节点对。

代码实现：

Copy
class Solution {
public:
vector<int> countPairsOfConnectableServers(vector<vector<int>> &edges, int signalSpeed) {
int n = edges.size() + 1; // 边比节点数少 1
// {root, [<child1, cost1>, <child2, cost2>, ..., <childn, costn>]}
vector<vector<pair<int, int>>> graph(n);

for (const auto &edge : edges) {
graph[edge[0]].push_back({edge[1], edge[2]});
graph[edge[1]].push_back({edge[1], edge[2]});
}

/// @brief
/// @param root 根节点，传入该参数是为了防止子树根节点走向该节点
/// @param child 子节点
/// @param remainder root 到 child 的根节点的距离
/// @return 返回以 root 根节点，child 为其中一个子树，子树有多少个满足到达路径可以被 signalSpeed 整除的节点
function<int(int, int, int)> dfs = [&](int root, int child, int remainder) -> int {
int res = 0;
if (remainder == 0) {
++res;
}

// 遍历子节点的所有子树
for (const auto &[j, cost] : graph[child]) {
// 排除根节点路径
if (j != root) {
res += dfs(child, j, (remainder + cost) % signalSpeed);
}
}
return res;
};

vector<int> res(n);

// 遍历根节点
for (int i = 0; i < n; ++i) {
int pre = 0; // 统计根节点子树的可连接对数

// 遍历根节点的所有子树
for (auto &[j, cost] : graph[i]) {
// 当前子树满足的条件的节点个数
int cnt = dfs(i, j, cost % signalSpeed);
res[i] += pre * cnt;
pre += cnt;
}
}
return res;
}
};