【点分治】luoguP2664 树上游戏

应该是一道中等难度的点分？麻烦在一些细节。

题目描述

lrb有一棵树，树的每个节点有个颜色。给一个长度为n的颜色序列，定义s(i,j) 为i 到j 的颜色数量。以及

现在他想让你求出所有的sum[i]

输入输出格式

输入格式：

第一行为一个整数n，表示树节点的数量

第二行为n个整数，分别表示n个节点的颜色c[1],c[2]……c[n]

接下来n-1行，每行为两个整数x,y，表示x和y之间有一条边

输出格式：

输出n行，第i行为sum[i]

说明

sum[1]=s(1,1)+s(1,2)+s(1,3)+s(1,4)+s(1,5)=1+2+3+2+2=10
sum[2]=s(2,1)+s(2,2)+s(2,3)+s(2,4)+s(2,5)=2+1+2+1+3=9
sum[3]=s(3,1)+s(3,2)+s(3,3)+s(3,4)+s(3,5)=3+2+1+2+3=11
sum[4]=s(4,1)+s(4,2)+s(4,3)+s(4,4)+s(4,5)=2+1+2+1+3=9
sum[5]=s(5,1)+s(5,2)+s(5,3)+s(5,4)+s(5,5)=2+3+3+3+1=12

对于40%的数据，n<=2000

对于100%的数据，1<=n,c[i]<=10^5

---

题目分析

想法一：按颜色拆贡献

这里应该是有一种小颜色大颜色的分块套路的。但是这个想法我只能解决全局路径的数量和，并不会落实到点的询问。

想法二：点分治

目前尚未归结出点分治适用的具体问题范围……不过这一题是可以用点分解决的。

考虑每一层点分树，我们只需要对它的节点处理贡献。这里的贡献分为两部分：重心答案；经过重心的路径对子树的贡献。

重心的答案只需要以它自身为根，遍历一边该层点分树即可。子树内的答案处理要略微麻烦一些，需要分颜色来考虑贡献。记$colCnt[i]$为所有以重心为起点的路径中，含有颜色$i$的路径条数。然后首先假定子树内所有点的答案都为$\sum colCnt[i]$，再容斥考虑重心到子树路径上的颜色所产生的贡献。

记当前点分树中除去正在处理的子树的大小为$outTot$，那么对于子树内点$x$，由于它具有颜色$c[x]$，所以对自身的答案有一个$outTot-colCnt[c[x]]$的贡献。并且，这一个贡献对于$x$的子树也是一概适用的，所以这一个标记要差分式地下传。

整体思路就是这些。这一题的点分涉及到例如“子树结构的重定向”或是“两个颜色桶并存”的一些细节问题，所以实现上面可能有一定的难度。

（话说这题的码风怎么这么丑）

 

#include<bits/stdc++.h>
typedef long long ll;
const int maxn = 100035;
const int maxm = 200035;

ll ans[maxn];
int n,bloTot,outTot,c[maxn]; 
int size[maxn],son[maxn],root;
int edgeTot,head[maxn],nxt[maxm],edges[maxm];
int cols,cnt,cl,col[maxn],colTmp[maxn],colTim[maxn],colCnt[maxn],subCnt[maxn];
bool colEx[maxn],divEx[maxn];

int read()
{
    char ch = getchar();
    int num = 0, fl = 1;
    for (; !isdigit(ch); ch=getchar())
        if (ch=='-') fl = -1;
    for (; isdigit(ch); ch=getchar())
        num = (num<<1)+(num<<3)+ch-48;
    return num*fl;
}
void addedge(int u, int v)
{
    edges[++edgeTot] = v, nxt[edgeTot] = head[u], head[u] = edgeTot;
    edges[++edgeTot] = u, nxt[edgeTot] = head[v], head[v] = edgeTot;
}
void getRoot(int x, int fa)
{
    size[x] = 1, son[x] = 0;
    for (int i=head[x]; i!=-1; i=nxt[i])
    {
        int v = edges[i];
        if (divEx[v]||v==fa) continue;
        getRoot(v, x), size[x] += size[v];
        son[x] = std::max(son[x], size[v]);
    }
    son[x] = std::max(son[x], bloTot-size[x]);
    if (son[x] < son[root]) root = x;
}
void colDfs(int x, int fa, int *cnt)
{
    if (!colEx[c[x]]) colEx[c[x]] = 1, col[++cols] = c[x];
    if ((++colTim[c[x]])==1) cnt[c[x]] += size[x];
    for (int i=head[x]; i!=-1; i=nxt[i])
        if ((!divEx[edges[i]])&&(edges[i]!=fa))
            colDfs(edges[i], x, cnt);
    --colTim[c[x]];
}
void colClear()
{
    for (int i=1; i<=cl; i++) colEx[colTmp[i]] = 0;
    cols = 0;
}
void modify(int x, int fa, ll tag)
{
    if ((++colTim[c[x]])==1) tag += outTot-colCnt[c[x]];
    ans[x] += tag+cnt;
    for (int i=head[x]; i!=-1; i=nxt[i])
    {
        int v = edges[i];
        if (v==fa||divEx[v]) continue;
        modify(v, x, tag);
    }
    --colTim[c[x]];
}
void calc(int rt)　　　　　　//核心操作在这里
{
    colClear(), getRoot(rt, 0);
    colDfs(rt, 0, colCnt);
    cnt = 0, cl = cols;
    for (int i=1; i<=cols; i++)
        cnt += colCnt[col[i]], colTmp[i] = col[i];
    ans[rt] += cnt;
    for (int i=head[rt]; i!=-1; i=nxt[i])
    {
        int v = edges[i];
        if (divEx[v]) continue;
        for (int j=1; j<=cl; j++) subCnt[colTmp[j]] = 0;　　//及时清除数组
        colClear();
        colEx[c[rt]] = 1;
        colDfs(v, rt, subCnt);　　　　　　　　　　//统计子树内的含颜色i路径条数
        colEx[c[rt]] = 0;
        colCnt[c[rt]] -= size[v], cnt -= size[v];　　　//除去重心出发的路径
        for (int j=1; j<=cols; j++)
        {
            colCnt[col[j]] -= subCnt[col[j]];　　//除去子树内的路径（因为考虑子树外路径）
            cnt -= subCnt[col[j]];
        }
        outTot = size[rt]-size[v], modify(v, rt, 0);　　//对子树内累加贡献
        colCnt[c[rt]] += size[v], cnt += size[v];　　　　//恢复处理子树前状态
        for (int j=1; j<=cols; j++)
        {
            colCnt[col[j]] += subCnt[col[j]];
            cnt += subCnt[col[j]];
        }
    }
    for (int i=1; i<=cl; i++)
        colCnt[colTmp[i]] = 0;　　　　　　//colTmp[]的作用；清空colCnt[]
}
void deal(int rt)
{
    calc(rt), divEx[rt] = 1;
    for (int i=head[rt]; i!=-1; i=nxt[i])
    {
        int v = edges[i];
        if (divEx[v]) continue;
        root = 0, bloTot = size[v];
        getRoot(v, 0), deal(root);
    }
}
int main()
{
    memset(head, -1, sizeof head);
    n = read(), son[0] = n;
    for (int i=1; i<=n; i++) c[i] = read();
    for (int i=1; i<n; i++) addedge(read(), read());
    bloTot = n, getRoot(1, 0), deal(root);
    for (int i=1; i<=n; i++) printf("%lld\n",ans[i]);
    return 0;
}

 

 

 

END