110 lines
3.0 KiB
C++
110 lines
3.0 KiB
C++
#include <iostream>
|
||
#include <vector>
|
||
#include <algorithm>
|
||
|
||
using namespace std;
|
||
|
||
// 建议使用 N 定义数组，比 vector 稍快且容易管理
|
||
const int N = 100005;
|
||
vector<pair<int, int>> adj[N];
|
||
bool vis[N]; // 分治标记：记录哪些点已经作为重心被“删掉”了
|
||
int treesize[N];
|
||
int root, min_maxpart, K;
|
||
long long ans = 0;
|
||
|
||
// 存储当前子树所有节点到重心的距离
|
||
vector<int> dists;
|
||
|
||
// 1. 找重心：注意这里不再需要 visited 数组，用 fa 限制方向即可
|
||
void get_root(int u, int fa, int currsize) {
|
||
treesize[u] = 1;
|
||
int maxpart = 0;
|
||
for (auto& edge : adj[u]) {
|
||
int v = edge.first;
|
||
if (v == fa || vis[v]) continue; // 关键：不能走已经“删掉”的点
|
||
get_root(v, u, currsize);
|
||
treesize[u] += treesize[v];
|
||
maxpart = max(maxpart, treesize[v]);
|
||
}
|
||
maxpart = max(maxpart, currsize - treesize[u]);
|
||
if (maxpart < min_maxpart) {
|
||
min_maxpart = maxpart;
|
||
root = u;
|
||
}
|
||
}
|
||
|
||
// 2. 收集距离：把所有节点到重心的距离存入 dists 数组
|
||
void get_dists(int u, int fa, int d) {
|
||
dists.push_back(d);
|
||
for (auto& edge : adj[u]) {
|
||
int v = edge.first;
|
||
int w = edge.second;
|
||
if (v == fa || vis[v]) continue;
|
||
get_dists(v, u, d + w);
|
||
}
|
||
}
|
||
|
||
// 3. 计算贡献：双指针法统计满足 dist1 + dist2 <= K 的对数
|
||
long long calc(int u, int init_dist) {
|
||
dists.clear();
|
||
get_dists(u, -1, init_dist);
|
||
sort(dists.begin(), dists.end());
|
||
|
||
long long count = 0;
|
||
int l = 0, r = dists.size() - 1;
|
||
while (l < r) {
|
||
if (dists[l] + dists[r] <= K) {
|
||
count += (r - l);
|
||
l++;
|
||
} else {
|
||
r--;
|
||
}
|
||
}
|
||
return count;
|
||
}
|
||
|
||
// 4. 点分治主函数
|
||
void divide(int u, int currsize) {
|
||
// 每次进入先找当前连通块的重心
|
||
min_maxpart = currsize + 7; // 初始化为一个大值
|
||
get_root(u, -1, currsize);
|
||
|
||
int r = root; // 锁定重心
|
||
vis[r] = true; // 【关键】逻辑切断重心
|
||
|
||
// 统计经过重心的路径
|
||
ans += calc(r, 0);
|
||
|
||
// 递归子树
|
||
for (auto& edge : adj[r]) {
|
||
int v = edge.first;
|
||
int w = edge.second;
|
||
if (vis[v]) continue;
|
||
|
||
// 【关键】去重：减去在同一个儿子子树内提前满足条件的路径
|
||
ans -= calc(v, w);
|
||
|
||
// 这里的 currsize 要更新为子树的实际大小
|
||
// 注意：由于已经跑过 get_root，treesize[v] 此时就是正确的
|
||
// 但如果 v 是 r 的“上方”节点，size 应该是 currsize - treesize[r]
|
||
int next_size = (treesize[v] < treesize[r]) ? treesize[v] : (currsize - treesize[r]);
|
||
divide(v, next_size);
|
||
}
|
||
}
|
||
|
||
int main() {
|
||
int n;
|
||
if (!(cin >> n >> K)) return 0;
|
||
for (int i = 0; i < n - 1; i++) {
|
||
int u, v, w;
|
||
cin >> u >> v >> w;
|
||
adj[u].push_back({v, w});
|
||
adj[v].push_back({u, w});
|
||
}
|
||
|
||
divide(1, n); // 从节点1开始分治，总大小为 n
|
||
|
||
cout << ans << endl;
|
||
return 0;
|
||
}
|