GFOJ510树上路径

Problem 510: 树上路径
Time Limit: 3000 ms
Memory Limit: 512000 KB

Problem Description

练一练点分……
dep表示点到重心距离，$_mx$表示点到重心最大值，对于$_mx$排序后乱搞……
表示拍了很久发现犯了个**错误……

#include <cstdio>
#include <algorithm>
#define maxn 100000+10
using namespace std;
int num,x,y,dep[maxn],_mx[maxn],rot,ans,now_sz,vis[maxn],sz[maxn],f[maxn],g[10000000+10],n,p,head[maxn],m,a[maxn];
struct xx
{
    int v,next;
}b[maxn<<1];
struct yy
{
    int dep,mx;
}c[maxn];
void add(int u,int v)
{
    b[++m]=(xx){v,head[u]};
    head[u]=m;
}

void getrot(int x,int fa)
{
    f[x]=0;sz[x]=1;
    for (int i=head[x];i;i=b[i].next)
    {
        int v=b[i].v;
        if (v==fa||vis[v]) continue;
        getrot(v,x);
        sz[x]+=sz[v];
        f[x]=max(f[x],sz[v]);
    }
    f[x]=max(now_sz-sz[x],f[x]);
    rot=f[x]<f[rot]?x:rot;
}

void dfs(int x,int fa)
{
    for (int i=head[x];i;i=b[i].next)
    {
        int v=b[i].v;
        if (vis[v]||v==fa)continue;
        /*dep[v]=(dep[x]+a[x])%p;
        _mx[v]=max(_mx[v],a[x]);*/

        dep[v]=(dep[x]+a[v])%p;
        _mx[v]=max(_mx[x],a[v]);
        c[++num]=(yy){dep[v],_mx[v]};
        dfs(v,x);
    }

}

bool cmp(const yy&a,const yy&b)
{
    return a.mx<b.mx;
}

int cal(int x,int f)
{
    int tmp=0;
    num=0;
    dep[x]=f==-1?a[x]%p:(a[x]+f)%p;
    _mx[x]=f==-1?a[x]:max(a[x],f);
    c[++num]=(yy){dep[x],_mx[x]};
    dfs(x,0);
    sort(c+1,c+num+1,cmp);
    for (int i=1;i<=num;++i)
    {
        int hh=(p+c[i].mx+(f==-1?a[x]:f)-c[i].dep)%p;//(f==-1?a[x]:f)不要写成a[x]!!! 
        tmp+=g[hh];
        ++g[c[i].dep];
    }
    for (int i=1;i<=num;++i) --g[c[i].dep];
    return tmp;
}


void work(int x)
{
    vis[x]=1;
    ans+=cal(x,-1);
    for (int i=head[x];i;i=b[i].next)
    {
        int v=b[i].v;
        if (vis[v])continue;
        ans-=cal(v,a[x]);
        now_sz=sz[v],rot=0;
        getrot(v,0);
        work(rot);
    }
}

int main()
{
    freopen("data.out","r",stdin);
    freopen("wa.out","w",stdout);
    scanf("%d%d",&n,&p);
    for (int i=1;i<n;++i) scanf("%d%d",&x,&y),add(x,y),add(y,x);
    for (int i=1;i<=n;++i) scanf("%d",&a[i]);
    f[0]=999999999;
    now_sz=n;
    getrot(1,0);
    work(rot);
    ans+=n;
    printf("%d",ans);
    return 0;
}