[LuoGu]P2664 树上游戏

Portal

这题真的好.

看到树上路径, 脑子里就要点分治

这一题对于每个点都要计算一遍, 如果暴算实在不好算, 这样我们就可以考虑算贡献.

直接计算每种颜色的贡献.

因为一条过重心的路径中, 可能两边都会有相同颜色, 那么我们就只考虑当前点到分治中心的链上.

然后把多算的这条链上的东西删除就可以了.

Copy From Here

那么我们就可以这样做了：

１.对树进行第一遍dfs，预处理size和上方性质中的贡献。（开一个color数组即可记录贡献），同时记录贡献总和sum

２.枚举跟的所有子树，先把子树扫一遍清除其在color数组中的所有贡献。接着，对于该子树中的每一个点ｊ：

设Ｘ＝sigma color[j 到根上（不包括根）的所有颜色] 由于这些颜色已经出现过，我们不能在该子树外计算其贡献）

设num为ｊ到根上（不包括根）的颜色数

设Ｙ为size[root]-size[该子树]（即所有其他子树＋根的点数）

则ans[j]+＝sum-X+num*Y;

３.别忘了计算单独ｒｏｏｔ的ａｎｓ

　　ans[root]+=sum-color[根的颜色]+size[root]

４．清空贡献数组以及其他东西

Code

#include<bits/stdc++.h>
using namespace std;
#define rep(i, a, b) for(int i = (a), i##_end_ = (b); i <= i##_end_; ++i)
#define drep(i, a, b) for(int i = (a), i##_end_ = (b); i >= i##_end_; --i)
#define clar(a, b) memset((a), (b), sizeof(a))
#define debug(...) fprintf(stderr, __VA_ARGS__)
typedef long long LL;
typedef long double LD;
int read() {
    char ch = getchar();
    int x = 0, flag = 1;
    for (;!isdigit(ch); ch = getchar()) if (ch == '-') flag *= -1;
    for (;isdigit(ch); ch = getchar()) x = x * 10 + ch - 48;
    return x * flag;
}
void write(int x) {
    if (x < 0) putchar('-'), x = -x;
    if (x >= 10) write(x / 10);
    putchar(x % 10 + 48);
}

const int Maxn = 100009;

struct edge {
    int to, nxt;
}g[Maxn << 1];

int n, a[Maxn], head[Maxn], e;

void add(int u, int v) {
    g[++e] = (edge){v, head[u]}, head[u] = e;
}

void init() {
    clar(head, -1);
    n = read();
    rep (i, 1, n) a[i] = read();
    rep (i, 1, n - 1) {
        int u = read(), v = read();
        add(u, v), add(v, u);
    }
}

LL ans[Maxn], vis[Maxn], Tmp, heart;
int size[Maxn];

void dfs(int u, int pa) { 
    size[u] = 1;
    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (v != pa && !vis[v]) {
            dfs(v, u);
            size[u] += size[v];
        }
    }
}

void centroid(int u, int pa, int field) {
    int bbs = field - size[u];
    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (!vis[v] && v != pa) {
            centroid(v, u, field);
            bbs = max(bbs, size[v]);
        }
    }
    if (bbs < Tmp) Tmp = bbs, heart = u;
}

LL color[Maxn], Sum, cnt[Maxn], size1[Maxn];

void dfs_init(int u, int pa) {
    size1[u] = 1; ++cnt[a[u]];
    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (v != pa && !vis[v]) {
            dfs_init(v, u);
            size1[u] += size1[v];
        }
    }

    if (cnt[a[u]] == 1) {
        color[a[u]] += size1[u];
        Sum += size1[u];
    }
    --cnt[a[u]];
}

void change(int u, int pa, int delta) {
    ++cnt[a[u]];
    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (v != pa && !vis[v]) change(v, u, delta);
    }

    if (cnt[a[u]] == 1) {
        color[a[u]] += size1[u] * 1ll * delta;
        Sum += size1[u] * 1ll * delta;
    }
    --cnt[a[u]];
}

void count(int u, int pa, int X, int Y, int num) {
    ++cnt[a[u]];
    if (cnt[a[u]] == 1) {
        X += color[a[u]];
        ++Y;
    }
    ans[u] += Sum - X + num * Y;
    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (!vis[v] && v != pa) count(v, u, X, Y, num);
    }
    --cnt[a[u]];
}

void dfs_clear(int u, int pa) {
    color[a[u]] = 0, cnt[a[u]] = 0, size1[u] = 0;
    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (v != pa && !vis[v]) dfs_clear(v, u);
    }
}

void calc(int u) {
    dfs_init(u, 0);
    ans[u] += Sum - color[a[u]] + size1[u];

    for (int i = head[u]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (!vis[v]) {
            ++cnt[a[u]];
            Sum -= size1[v]; color[a[u]] -= size1[v];   
            change(v, u, -1);
            --cnt[a[u]];
            count(v, u, 0, 0, size1[u] - size1[v]);
            ++cnt[a[u]];
            Sum += size1[v]; color[a[u]] += size1[v];   
            change(v, u, 1);
            --cnt[a[u]];
        }
    }

    dfs_clear(u, 0), Sum = 0;
}

void divConquer(int u) {
    dfs(u, 0), Tmp = INT_MAX;
    centroid(u, 0, size[u]);
    calc(heart); vis[heart] = 1;
    for (int i = head[heart]; ~i; i = g[i].nxt) {
        int v = g[i].to;
        if (!vis[v]) divConquer(v);
    }
}

void solve() {
    divConquer(1);
    rep (i, 1, n) printf("%lld\n", ans[i]);
}

int main() {
    freopen("LG2664.in", "r", stdin);
    freopen("LG2664.out", "w", stdout);

    init();
    solve();

#ifdef Qrsikno
    debug("\nRunning time: %.3lf(s)\n", clock() * 1.0 / CLOCKS_PER_SEC);
#endif
    return 0;
}

转载于:https://www.cnblogs.com/qrsikno/p/10247197.html