相比两个点的距离相加等于k, 等于k的倍数而言
这题稍微转换一下思路即可
利用容斥原理去掉重复计算的点对
点击查看代码
#include <iostream> #include <cstdio> #include <cstring> #include <algorithm> using namespace std; #define IOS ios::sync_with_stdio(false), cin.tie(0), cout.tie(0) using ll = long long; constexpr int MAXN = 4e4 + 7; int h[MAXN], e[MAXN << 1], ne[MAXN << 1], w[MAXN << 1], Size[MAXN], mx[MAXN], idx, root, S, subSize; int ans, dis[MAXN], tot; int dd[MAXN]; bool vis[MAXN]; //当前点是否被删除 int n, m, k, cnt; void add(int u, int v, int c = 0) { ++idx; e[idx] = v; w[idx] = c; ne[idx] = h[u]; h[u] = idx; } void getroot(int x, int fa) //以下是树的重心模板 { Size[x] = 1; mx[x] = 0; for (int i = h[x]; i; i = ne[i]) { int j = e[i]; if (j != fa && !vis[j]) { getroot(j, x); Size[x] += Size[j]; mx[x] = max(mx[x], Size[j]); } } mx[x] = max(mx[x], S - Size[x]); if (mx[x] < mx[root]) root = x; } void getdis(int x, int fa) { dd[++cnt] = dis[x]; for (int i = h[x]; i; i = ne[i]) { int j = e[i]; if (j != fa && !vis[j]) { dis[j] = dis[x] + w[i]; getdis(j, x); //遍历以j节点为根节点的整颗子树 } } } int solve(int x) { int ans = 0; cnt = 0; //将上次保存点的距离的计数清空( getdis(x, 0); //以x为根统计一遍所有子节点的距离 sort(dd + 1, dd + cnt + 1); //排序后用双指针从头从尾开始尺取 for (int u = 1, v = cnt; u <= v;) { if (dd[u] + dd[v] <= k) { ans += v - u, ++u; } else --v; } return ans; } void div(int x) { vis[x] = true; //删除当前根节点 dis[x] = 0; ans += solve(x); //路径计数! for (int i = h[x]; i; i = ne[i]) { int j = e[i]; if (!vis[j]) // j节点被删除则不搜索, 因为已经被搜索过 { dis[j] = w[i]; ans -= solve(j); root = 0; S = Size[j]; //供找重心使用 mx[0] = Size[j]; //将最重的儿子大小暂时设为以j为根节点的子树大小, 留给后面找重心使用 getroot(j, 0); getroot(root, 0); div(root); //找到重心, 将重心作为根节点开始搜子树(保证时间复杂度最优) } } } int main() { int u, v, c; IOS; cin >> n; for (int i = 1; i < n; ++i) { cin >> u >> v >> c; add(u, v, c); add(v, u, c); } cin >> k; root = 0; mx[0] = n; S = n; getroot(1, 0); getroot(root, 0); div(root); cout << ans << '\n'; return 0; }