洛谷 P2634 聪聪可可 —— 树形DP / 点分治

题目：https://www.luogu.org/problemnew/show/P2634

今天刚学了点分治，做例题；

好不容易A了，结果发现自己写的是树形DP...(也不用找重心)(比点分治快)

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
int const maxn=20005,inf=0x3f3f3f3f;
int n,hd[maxn],ct,siz[maxn],mx,rt;
ll cnt[maxn][5],ans,tot,tmp[5];
struct N{
    int to,nxt,w;
    N(int t=0,int n=0,int w=0):to(t),nxt(n),w(w) {}
}ed[maxn<<1];
void add(int x,int y,int z){ed[++ct]=N(y,hd[x],z); hd[x]=ct;}
ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
void dfs(int x,int fa)
{
    siz[x]=1; int nmx=0;
    for(int i=hd[x],u;i;i=ed[i].nxt)
    {
        if((u=ed[i].to)==fa)continue;
        dfs(u,x); siz[x]+=siz[u]; 
        nmx=max(nmx,siz[u]); 
    }
    nmx=max(nmx,n-siz[x]);
    if(nmx<mx)rt=x,mx=nmx;//mx=nmx!!
}
void work(int x,int fa)
{
    cnt[x][0]=1;
    for(int i=hd[x],u;i;i=ed[i].nxt)
    {
        if((u=ed[i].to)==fa)continue;
        work(u,x);
        for(int j=0;j<=2;j++)tmp[j]=cnt[u][((j-ed[i].w)%3+3)%3];
        ans+=cnt[x][1]*tmp[2]; 
        ans+=cnt[x][2]*tmp[1]; 
        ans+=cnt[x][0]*tmp[0]; 
        for(int j=0;j<=2;j++)cnt[x][j]+=tmp[j];
    }
//    ans+=cnt[x][0];//
//    ans++;
}
int main()
{
//    freopen("test.txt","r",stdin);
    scanf("%d",&n); mx=inf;
    for(int i=1,x,y,z;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z); add(y,x,z);
    }
    dfs(1,0); work(rt,0); 
    ans=ans*2+n; tot=(ll)n*n; 
    ll g=gcd(ans,tot);
    printf("%lld/%lld
",ans/g,tot/g);
    return 0;
}

树形DP

 点分治是先算出经过根的，再去掉根，对于子树进行相同的操作，还要容斥掉刚才算上的、经过该子树根节点的答案。

代码如下：

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
int const maxn=20005,inf=0x3f3f3f3f;
int n,hd[maxn],ct,siz[maxn],mx,rt,sum,dep[maxn];
ll cnt[maxn][5],ans,tot,tmp[5];
bool vis[maxn];
struct N{
    int to,nxt,w;
    N(int t=0,int n=0,int w=0):to(t),nxt(n),w(w) {}
}ed[maxn<<1];
void add(int x,int y,int z){ed[++ct]=N(y,hd[x],z); hd[x]=ct;}
ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
void dfs(int x,int fa)
{
    siz[x]=1; int nmx=0;
    for(int i=hd[x],u;i;i=ed[i].nxt)
    {
        if((u=ed[i].to)==fa||vis[u])continue;
        dfs(u,x); siz[x]+=siz[u]; 
        nmx=max(nmx,siz[u]); 
    }
    nmx=max(nmx,sum-siz[x]);
    if(nmx<mx)rt=x,mx=nmx;//mx=nmx!!
}
void getdep(int x,int fa)
{
    tmp[dep[x]%=3]++;
    for(int i=hd[x],u;i;i=ed[i].nxt)
    {
        if((u=ed[i].to)==fa||vis[u])continue;
        dep[u]=dep[x]+ed[i].w;
        getdep(u,x);
    }
}
ll calc(int x,int v)
{
    dep[x]=v; tmp[0]=tmp[1]=tmp[2]=0;
    getdep(x,0);//0 
    return tmp[0]*tmp[0]+tmp[1]*tmp[2]*2;
}
void work(int x)
{
    ans+=calc(x,0); vis[x]=1;
    for(int i=hd[x],u;i;i=ed[i].nxt)
    {
        if(vis[u=ed[i].to])continue;
        ans-=calc(u,ed[i].w); 
        mx=inf; sum=siz[u]; dfs(u,0);
        work(rt);
    }
}
int main()
{
//    freopen("test.txt","r",stdin);
    scanf("%d",&n);
    for(int i=1,x,y,z;i<n;i++)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z); add(y,x,z);
    }
    mx=inf; sum=n; dfs(1,0); work(rt); 
    tot=(ll)n*n; 
    ll g=gcd(ans,tot);
    printf("%lld/%lld
",ans/g,tot/g);
    return 0;
}