[COCI2020-2021#4] Janjetina 题解

· · 题解

题目链接

题目大意:

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

分析:

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

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

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

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

g[u]-d[u]-d[v]\geq k,& g[u]\geq g[v]\\ g[v]-d[u]-d[v]\geq k,& g[u]\leq g[v] \end{cases}

移项后:

g[u]-d[u]-k\geq d[v],& g[u]\geq g[v]\\ -d[u]-k\geq d[v]-g[v],& g[u]\leq g[v] \end{cases}

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

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

时间复杂度分析:

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

代码:

#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;
}