AcWing252. 树
题意
给定一棵 n n n 个节点的树,每条边有边权,求出树上两点距离小于等于 k k k 的点对数量。
(本题与洛谷 P4178 Tree 非常相似,但是这题 k k k 比较大,且存在零边,当然也可以和洛谷那题一样用树状数组+特判零边来完成,这里主要讲一下容斥的方法。)
解法
- 考虑以
u
u
u 作为路径的
lca
, u u u 连接了若干子树,遍历每一棵子树,计算出深度。将所有的深度排序(包括根),然后用双指针来计算有多少对距离之和小于等于 k k k 。这样计算的包含一些不合法的路径,可以发现当两个点来自不同的子树的时候一定是合法的,不合法来自于两个点在用一个子树上,这样lca
就不对了。所以再减去同一棵子树的贡献即可。 - 每次枚举的
lca
都应该是树的重心,否则可以在链上会退化成 O ( n 2 ) O(n^2) O(n2) 。
代码
容斥:
#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e4 + 5;
int n, k;
vector<pair<int, int>> g[maxn];
int rt, d[maxn], sz[maxn];
bool vis[maxn];
void dfs_rt(int u, int f, int tot) {
sz[u] = 1;
int maxx = 0;
for (auto e : g[u]) {
int v = e.first;
if (vis[v] || v == f) continue;
dfs_rt(v, u, tot);
sz[u] += sz[v];
maxx = max(maxx, sz[v]);
}
maxx = max(maxx, tot - sz[u]);
if (2 * maxx <= tot) rt = u;
}
int cnt, tmp[maxn];
void dfs_d(int u, int f) {
tmp[++cnt] = d[u];
for (auto e : g[u]) {
int v = e.first, w = e.second;
if (vis[v] || v == f) continue;
d[v] = d[u] + w;
dfs_d(v, u);
}
}
int cul(int u, int f) {
cnt = 0;
dfs_d(u, f);
sort(tmp + 1, tmp + 1 + cnt);
int l = 1, r = cnt, ans = 0;
while (l < r) {
while (l < r - 1 && tmp[l] + tmp[r] > k) --r;
if (tmp[l] + tmp[r] > k) break;
ans += r - l;
++l;
}
return ans;
}
int work(int u, int f, int tot) {
dfs_rt(u, f, tot);
u = rt, vis[u] = 1, d[u] = 0;
int ans = cul(u, f);
for (auto e : g[u]) {
int v = e.first;
if (vis[v]) continue;
ans -= cul(v, u);
ans += work(v, u, cnt);
}
return ans;
}
int main() {
while (scanf("%d%d", &n, &k) && n && k) {
rep(i, 1, n) g[i].clear(), vis[i] = 0;
rep(i, 1, n - 1) {
int u, v, w;
scanf("%d%d%d", &u, &v, &w);
++u, ++v;
g[u].push_back({v, w});
g[v].push_back({u, w});
}
printf("%d\n", work(1, 0, n));
}
}
树状数组:
#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e4 + 5;
int n, k;
vector<pair<int, int>> g[maxn];
int sz[maxn], rt, d[maxn];
bool vis[maxn];
int c[int(5e6 + 5)];
void add(int pos, int x) {
while (pos <= k) c[pos] += x, pos += pos & -pos;
}
int getsum(int pos) {
int ans = 0;
while (pos) ans += c[pos], pos -= pos & -pos;
return ans + c[0];
}
void init() {
rep(i, 1, n) g[i].clear(), vis[i] = 0;
}
void dfs_rt(int u, int f, int tot) {
sz[u] = 1;
int maxx = 0;
for (auto e : g[u]) {
int v = e.first;
if (vis[v] || v == f) continue;
dfs_rt(v, u, tot);
sz[u] += sz[v];
maxx = max(maxx, sz[v]);
}
maxx = max(maxx, tot - sz[u]);
if (maxx * 2 <= tot) rt = u;
}
void dfs_ans(int u, int f, int &ans, int &cnt) {
++cnt;
if (d[u] <= k) ans += getsum(k - d[u]) + 1;
for (auto e : g[u]) {
int v = e.first, w = e.second;
if (vis[v] || v == f) continue;
d[v] = d[u] + w;
dfs_ans(v, u, ans, cnt);
}
}
void dfs_update(int u, int f, int val) {
if (d[u] <= k) {
if (!d[u])
c[d[u]] += val;
else
add(d[u], val);
}
for (auto e : g[u]) {
int v = e.first;
if (vis[v] || v == f) continue;
dfs_update(v, u, val);
}
}
int work(int u, int f, int tot) {
dfs_rt(u, f, tot);
u = rt, vis[u] = 1, d[u] = 0;
int ans = 0;
for (auto e : g[u]) {
int v = e.first, w = e.second, cnt;
if (vis[v]) continue;
d[v] = w, cnt = 0;
dfs_ans(v, u, ans, cnt);
sz[v] = cnt;
dfs_update(v, u, 1);
}
for (auto e : g[u]) {
int v = e.first;
if (vis[v]) continue;
dfs_update(v, u, -1);
}
for (auto e : g[u]) {
int v = e.first;
if (vis[v]) continue;
ans += work(v, u, sz[v]);
}
return ans;
}
int main() {
while (scanf("%d%d", &n, &k) && n && k) {
init();
rep(i, 1, n - 1) {
int u, v, w;
scanf("%d%d%d", &u, &v, &w), ++u, ++v;
g[u].push_back({v, w});
g[v].push_back({u, w});
}
printf("%d\n", work(1, 0, n));
}
}