[luogu2664]树上游戏

前言

点分治

题目相关

链接

题目大意

给出一棵树，每个节点有一个颜色
设$s(u,v)$为树上点$u$到点$v$的路径上的颜色数
对于每个$i$，求$$S_i=\sum _{j=1}^{n}s(i,j)$$

数据范围

$n\le 100000$

题解

点分治大法好（但是我们这里不考虑）
考虑每种颜色的贡献
我们对于一种颜色把所有点拿出来
我们发现这些点将整棵树切成了若干联通块（即不是这个颜色的点的联通快）
我们发现对于在同一块内的两个点为端点的路径是不包含这种颜色的，其它路径都包含
那么我们现在考虑用所有的边的方案减去不要求的方案
我们发现我们可以直接维护当前子树内每种颜色的边界节点，并用树上差分维护答案的区间减
最后dfs一遍下传差分标记即可
容易发现这些边界节点会被加入一次再删除一次
故总复杂度$\mathcal{O}(n)$

代码

#include<cstdio>
#include<cctype>
#include<algorithm>
#include<vector>
namespace fast_IO
{
    const int IN_LEN=1000000,OUT_LEN=1000000;
    char ibuf[IN_LEN],obuf[OUT_LEN],*ih=ibuf+IN_LEN,*oh=obuf,*lastin=ibuf+IN_LEN,*lastout=obuf+OUT_LEN-1;
    inline char getchar_(){return (ih==lastin)&&(lastin=(ih=ibuf)+fread(ibuf,1,IN_LEN,stdin),ih==lastin)?EOF:*ih++;}
    inline void putchar_(const char x){if(oh==lastout)fwrite(obuf,1,oh-obuf,stdout),oh=obuf;*oh++=x;}
    inline void flush(){fwrite(obuf,1,oh-obuf,stdout);}
}
using namespace fast_IO;
#define getchar() getchar_()
#define putchar(x) putchar_((x))
#define rg register
typedef long long ll;
template <typename T> inline T max(const T a,const T b){return a>b?a:b;}
template <typename T> inline T min(const T a,const T b){return a<b?a:b;}
template <typename T> inline void mind(T&a,const T b){a=a<b?a:b;}
template <typename T> inline void maxd(T&a,const T b){a=a>b?a:b;}
template <typename T> inline T abs(const T a){return a>0?a:-a;}
template <typename T> inline void Swap(T&a,T&b){T c=a;a=b;b=c;}
//template <typename T> inline void swap(T*a,T*b){T c=a;a=b;b=c;}
template <typename T> inline T gcd(const T a,const T b){if(!b)return a;return gcd(b,a%b);}
template <typename T> inline T lcm(const T a,const T b){return a/gcd(a,b)*b;}
template <typename T> inline T square(const T x){return x*x;};
template <typename T> inline void read(T&x)
{
    char cu=getchar();x=0;bool fla=0;
    while(!isdigit(cu)){if(cu=='-')fla=1;cu=getchar();}
    while(isdigit(cu))x=x*10+cu-'0',cu=getchar();
    if(fla)x=-x;
}
template <typename T> inline void printe(const T x)
{
    if(x>=10)printe(x/10);
    putchar(x%10+'0');
}
template <typename T> inline void print(const T x)
{
    if(x<0)putchar('-'),printe(-x);
    else printe(x);
}
const int maxn=100005,maxm=200005;
int n,c[maxn];
int head[maxn],nxt[maxm],tow[maxm],tmp;
inline void addb(const int u,const int v)
{
    tmp++;
    nxt[tmp]=head[u];
    head[u]=tmp;
    tow[tmp]=v;
}
int size[maxn];
ll ans[maxn];
std::vector<int>col[maxn];
void dfs(const int u,const int fa)
{
    std::vector<int>&me=col[c[u]];
    const int SZ=me.size();
    size[u]=1;
    for(rg int i=head[u];i;i=nxt[i])
    {
        const int v=tow[i];
        if(v==fa)continue;
        dfs(v,u),size[u]+=size[v];
        if(u==0)
        {
            for(c[u]=1;c[u]<=100000;c[u]++)
            {
                std::vector<int>&M=col[c[u]];
                int dq=size[v];
                for(int i=M.size()-1;i>=SZ;i--)dq-=size[M[i]];
                ans[v]-=dq;
                for(int i=M.size()-1;i>=SZ;i--)ans[M[i]]+=dq;
                while(M.size()>0)M.pop_back();
            }
            continue;
        }
        int dq=size[v];
        for(int i=me.size()-1;i>=SZ;i--)dq-=size[me[i]];
        ans[v]-=dq;
        for(int i=me.size()-1;i>=SZ;i--)ans[me[i]]+=dq;
        while(me.size()>SZ)me.pop_back();
    }
    me.push_back(u);
}
void cf(const int u,const int fa)
{
    for(rg int i=head[u];i;i=nxt[i])
    {
        const int v=tow[i];
        if(v==fa)continue;
        ans[v]+=ans[u],cf(v,u);
    }
}
int main()
{
    read(n);
    for(rg int i=1;i<=n;i++)read(c[i]);
    addb(0,1),addb(1,0);
    for(rg int i=1;i<n;i++)
    {
        int u,v;read(u),read(v);
        addb(u,v),addb(v,u);
    }
    dfs(0,0);
    cf(0,0);
    for(rg int i=1;i<=n;i++)print(ans[i]+(ll)n*100000),putchar('\n');
    return flush(),0;
}

总结

点分好难，还是差分清真