做题记录P2664 树上游戏

P2664 树上游戏

一道很好的点分治练手题，虽然对于初学者可能难了点。题解里的神仙解法看得我一脸懵逼

乍一看题目，似乎是跟树上路径有关，可以用点分治，其实我是学了点分治然后来做的这道题。

对于每一个分治中心，有两种情况

子树中的点对当前的这个点的影响

经过该点的路径对其他子树中的点的影响

对于第一种情况，很好求，

我们用 $cnt_i$ 表示路径中包含颜色 $i$ 的路径数量, $Sum$ 表示所有颜色的路径数量和, $c_i$ 表示每个点的数量, $siz_i$ 表示以 $i$ 为根的子树大小

那么对当前这个点 $i$ 答案的贡献就是 $Sum-cnt_{c_i}+siz_{i}$, 减去一个 $cnt_{c_i}$ 是为了去掉重复计算的.

对于第二种情况就比较麻烦了

我们在计算的时候不考虑子树内部对答案的贡献,之考虑其他子树对当前子树内的点的答案贡献,所以在我们求出来的 $cnt, Sum$ 中把当前子树的贡献先删掉

然后 $dfs$ 这个子树,对于我们当前枚举的结点 $x$, 如果 $c_x$ 是第一次出现,那么对于 $x$ 的所有儿子节点,都需要减去一个$cnt_{c_x}$,然后加上除去这棵子树的其它子树的结点数量,

因为对于其他子树的任意一个结点,都可以经过 $x$ 这个点使路径上包含颜色 $c_x$.

代码实现,结合注释理解

#include<iostream>
#include<cstdio>
using namespace std;
const int N = 2e5 + 5;
typedef long long ll;

int n, top, Y;
int head[N], ver[N], net[N], idx, Sum;
int c[N], cnt[N], tim[N], siz[N], cut[N];
bool st[N], use[N];
ll sum[N];

void add(int a, int b)
{
    net[++idx] = head[a], ver[idx] = b, head[a] = idx;
}

int get_wc(int u, int fa, int tot, int &wc)
{
    int alls = 1, ms = 0;
    for (int i = head[u]; i; i = net[i])
    {
        int v = ver[i];
        if (v == fa || st[v])
            continue;
        int t = get_wc(v, u, tot, wc);
        alls += t;
        ms = max(ms, t);
    }
    ms = max(ms, tot - alls);
    if (ms <= tot / 2)
        wc = u;
    return alls;
}//求重心,不用多讲了吧

void get_siz(int u, int fa)
{
    siz[u] = 1;
    for (int i = head[u]; i; i = net[i])
        if (ver[i] != fa && !st[ver[i]])
            get_siz(ver[i], u), siz[u] += siz[ver[i]];
}

void up_col(int u, int fa, int tp)
{
    if (!tim[c[u]])
    {
        cnt[c[u]] += tp * siz[u];//当前颜色的贡献
        Sum += tp * siz[u];//总贡献
    }
    if (!use[c[u]] && tp == 1)
        cut[++top] = c[u], use[c[u]] = true;//记录一下一共用了哪些颜色
    tim[c[u]]++;//次数加一
    for (int i = head[u]; i; i = net[i])
        if (ver[i] != fa && !st[ver[i]])
            up_col(ver[i], u, tp);
    tim[c[u]]--;//还原
}

void update(int u, int fa, int num, ll tot)//更新子树内每个点的答案
{
    if (!tim[c[u]])
        num++, tot += cnt[c[u]];
    tim[c[u]]++;
    sum[u] += Sum - tot + num * Y;
    for (int i = head[u]; i; i = net[i])
        if (ver[i] != fa && !st[ver[i]])
            update(ver[i], u, num, tot);
    tim[c[u]]--;
}

void calc(int u)
{
    get_wc(u, -1, siz[u], u);//找出重心
    st[u] = true, Sum = top = 0;
    get_siz(u, -1), up_col(u, -1, 1);//求出子树大小与 $cnt$
    for (int i = head[u]; i; i = net[i])
    {
        int v = ver[i];
        if (!st[v])
        {
            tim[c[u]]++, up_col(v, u, -1);
            cnt[c[u]] -= siz[v], Sum -= siz[v];//先减去当前子树对 $cnt,Sum$ 的贡献
            Y = siz[u] - siz[v], update(v, u, 0, 0);
            up_col(v, u, 1), tim[c[u]]--;
            cnt[c[u]] += siz[v], Sum += siz[v];//加回来
        }
    }
    sum[u] += Sum - cnt[c[u]] + siz[u];
    for (int i = 1; i <= top; i++)
        use[cut[i]] = false, cnt[cut[i]] = 0;//清空数组
    for (int i = head[u]; i; i = net[i])
        if (!st[ver[i]])
            calc(ver[i]);//进行下一层点分
}

int main()
{
    scanf("%d", &n);
    for (int i = 1; i <= n; i++)
        scanf("%d", &c[i]);
    for (int i = 1; i < n; i++)
    {
        int a, b;
        scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }
    get_siz(1, -1);//求出每个子树的大小
    calc(1);//点分治
    for (int i = 1; i <= n; i++)
        printf("%lld\n", sum[i]);
    return 0;
}