BZOJ-2152: 聪聪可可

Description

求树上所有链中长度为3的倍数的链的个数。n<=20000。

树形DP。一开始居然想写多个log的点分治。
num[x][y]表示以x为节点，在其子树中找另一端点，所形成的链的长度mod 3等于y的个数。随便推一下就行了。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
#define N 21000
int n,num[N][3],ans;
int head[N],to[N<<1],nxt[N<<1],data[N<<1],cnt;
void add(int x,int y,int z)
{
    to[++cnt]=y;
    nxt[cnt]=head[x];
    head[x]=cnt;
    data[cnt]=z%3;
}
void dfs(int x,int pre)
{
    num[x][0]++;
    for(int i=head[x];i;i=nxt[i])
        if(to[i]!=pre)
        {
            dfs(to[i],x);
            ans+=num[x][0]*num[to[i]][(3-data[i])%3];
            ans+=num[x][1]*num[to[i]][(2-data[i]+3)%3];
            ans+=num[x][2]*num[to[i]][(1-data[i]+3)%3];
            num[x][data[i]]+=num[to[i]][0];
            num[x][(data[i]+1)%3]+=num[to[i]][1];
            num[x][(data[i]+2)%3]+=num[to[i]][2];
        }
}
int gcd(int x,int y){return !y?x:gcd(y,x%y);}
int main()
{
    scanf("%d",&n);
    int x,y,z,i;
    for(i=1;i!=n;++i)
    {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z); add(y,x,z);
    }
    dfs(1,0);
    ans=ans*2+n;
    int tmp=gcd(ans,n*n);
    printf("%d/%d",ans/tmp,n*n/tmp);
    return 0;
}