[COCI2020-2021#4] Janjetina 题解

题目链接

题目大意:

给定一棵树，求满足路径最大值减路径长度大于等于 $k$ 的点对 $(u,v)$ 的数量。

分析:

求树上满足条件的点对数量，很容易想到点分治可以做。

设当前根为 $root$，$g[x]$ 表示 $x$ 到 $root$ 之间的最大值，$d[x]$ 表示 $x$ 到 $root$ 的距离。

在 $root$ 为根的子树中点对 $(u,v)$ 合法当且仅当 $max(g[u],g[v])-d[u]-d[v]\ge k$ 且 $u$，$v$ 不属于 $root$ 的同一直接儿子下。

max 很不好处理，所以我们可以将它拆开成两个部分。

$$\{\begin{array}{ll}g[u]-d[u]-d[v]\ge k,& g[u]\ge g[v]\\ g[v]-d[u]-d[v]\ge k,& g[u]\le g[v]\end{array}$$

移项后:

$$\{\begin{array}{ll}g[u]-d[u]-k\ge d[v],& g[u]\ge g[v]\\ -d[u]-k\ge d[v]-g[v],& g[u]\le g[v]\end{array}$$

这两个部分显然可以用两个树状数组维护，至于如何确定是哪个贡献，直接对子树节点按 $g$ 排序即可。注意此处需要取等。

然后我们就发现会算重，因为 $root$ 同一直接儿子内部可能导致不合法的贡献，直接减去就完了。。。

时间复杂度分析:

经典点分治，每次找重心处理，只有 $\log n$ 层，每次计算答案，对但前子树扫描一遍预处理，$O(siz(x))$，排序加树状数组 $O(siz(x)\log (siz(x)))$，去重时同样是遍历一遍当前子树复杂度相同。子树总和大小为 $n$，所以最后时间复杂度为 $O(n{\log }^{2}n)$。

代码:

#include <bits/stdc++.h>
using namespace std;
typedef long long i64;
i64 read() {
	i64 x(0), f(0); char ch = getchar();
	while (!isdigit(ch)) f |= (ch == '-'), ch = getchar();
	while (isdigit(ch)) x = x * 10 + ch - '0', ch = getchar();
	return f ? -x : x;
}
int __stk[128], __len;
void put(i64 x) {
	if (x < 0) putchar('-'), x = -x;
	do { __stk[++__len] = x % 10, x /= 10; } while (x);
	while (__len) putchar(__stk[__len--] ^ 48);
}
const int N = 1e5 + 10, inf = 1e9;
int n, k;
i64 ans;
namespace DFZ {
    int head[N], cur;
    struct edge {
        int to, nxt, w;
    } e[N << 1];
    void link(int u, int v) {
		int w = read();
        e[++cur] = (edge) {v, head[u], w}, head[u] = cur;
        e[++cur] = (edge) {u, head[v], w}, head[v] = cur;
    }
    int root, siz[N], maxp[N];
    bool vis[N];
    void get_root(int u, int fa, int total) {
        maxp[u] = 0, siz[u] = 1;
        for (int i = head[u]; i; i = e[i].nxt) {
            if (e[i].to == fa || vis[e[i].to]) continue;
            get_root(e[i].to, u, total);
            siz[u] += siz[e[i].to], maxp[u] = max(maxp[u], siz[e[i].to]);
        }
        maxp[u] = max(maxp[u], total - siz[u]);
        if (maxp[u] < maxp[root]) root = u;
    }
	int g[N], d[N], o[N], cnt;
	bool cmp(int x, int y) { return g[x] < g[y]; }
	struct BIT {
		int t[2000000];
		void add(int x, int y) { for (x += 1000000; x <= 2000000; x += x & -x) t[x] += y; }
		int ask(int x) { int res = 0; for (x += 1000000; x; x -= x & -x) res += t[x]; return res; }
	}t1, t2;
	void prepare(int u, int fa) {
		for (int i = head[u]; i; i = e[i].nxt) 
			if (!vis[e[i].to] && e[i].to != fa) 
				g[e[i].to] = max(g[u], e[i].w), d[e[i].to] = d[u] + 1, prepare(e[i].to, u);
	}
	void dfs(int u, int fa) {
		o[++cnt] = u;
		for (int i = head[u]; i; i = e[i].nxt) 
			if (!vis[e[i].to] && e[i].to != fa) dfs(e[i].to, u);
	}
    void calc(int u) {
        g[u] = 0, d[u] = 0, prepare(u, 0);
		cnt = 0, dfs(u, 0), sort(o + 1, o + cnt + 1, cmp);
		for (int i = 1; i <= cnt; ++i) ans += t1.ask(g[o[i]] - d[o[i]] - k), t1.add(d[o[i]], 1);
		for (int i = cnt; i >= 1; --i) ans += t2.ask(-d[o[i]] - k), t2.add(d[o[i]] - g[o[i]], 1);
		for (int i = 1; i <= cnt; ++i) t1.add(d[o[i]], -1), t2.add(d[o[i]] - g[o[i]], -1);
		for (int i = head[u]; i; i = e[i].nxt) {
			if (vis[e[i].to]) continue;
			cnt = 0, dfs(e[i].to, u), sort(o + 1, o + cnt + 1, cmp);
			for (int j = 1; j <= cnt; ++j) ans -= t1.ask(g[o[j]] - d[o[j]] - k), t1.add(d[o[j]], 1);
			for (int j = cnt; j >= 1; --j) ans -= t2.ask(-d[o[j]] - k), t2.add(d[o[j]] - g[o[j]], 1);
			for (int j = 1; j <= cnt; ++j) t1.add(d[o[j]], -1), t2.add(d[o[j]] - g[o[j]], -1);
		}
    }
    void divide(int u) {
        calc(u), get_root(u, 0, 0), vis[u] = 1;
        for (int i = head[u]; i; i = e[i].nxt) {
            if (vis[e[i].to]) continue;
            root = 0, get_root(e[i].to, 0, siz[e[i].to]);
            divide(root);
        }
    }
}

signed main() {
	n = read(), k = read(); 
	for (int i = 1; i < n; ++i) DFZ::link(read(), read());
	DFZ::maxp[0] = n, DFZ::get_root(1, 0, n);
	DFZ::divide(DFZ::root), put(ans), putchar('\n');
	return 0;
}