Furukawa Nagisa's Tree

#include<bits/stdc++.h>
#define reg register int
#define il inline
#define fi first
#define se second
#define mk(a,b) make_pair(a,b)
#define numb (ch^'0')
using namespace std;
typedef long long ll;
template<class T>il void rd(T &x){
    char ch;x=0;bool fl=false;
    while(!isdigit(ch=getchar()))(ch=='-')&&(fl=true);
    for(x=numb;isdigit(ch=getchar());x=x*10+numb);
    (fl==true)&&(x=-x);
}
template<class T>il void ot(T x){x/10?ot(x/10):putchar(x%10+'0');}
template<class T>il void prt(T a[],int st,int nd){for(reg i=st;i<=nd;++i) printf("%lld ",a[i]);putchar('\n');}

namespace Miracle{
const int N=1e5+5;
int n;
ll mod,k,M;
ll mi[N];
vector<int>to[N];
struct hash{
    map<int,int>mp;
    int sz;
    vector<int>ad;
    void clear(){
        mp.clear();sz=0;
    }
    void ins(int v){
        ++mp[v];++sz;
    }
    void push(int v){
        ad.push_back(v);
    }
    void upda(){
        for(reg i=0;i<(int)ad.size();++i){
            ++mp[ad[i]];++sz;
        }
        ad.clear();
    }
}in,out;
int in0[N],in1[N],out0[N],out1[N];
int sz[N];
int v[N];
int rt;
bool vis[N];
void dfs1(int x,int fa){
    sz[x]=1;
    int mxsz=0;
    for(reg i=0;i<(int)to[x].size();++i){
        int y=to[x][i];
        if(vis[y]||y==fa) continue;
        dfs1(y,x);
        sz[x]+=sz[y];
        mxsz=max(mxsz,sz[y]);
    }
    mxsz=max(mxsz,nowsz-sz[x]);
    if(mxsz<=nowsz/2){
        rt=x;
    }
}
void dfs2(int x,int fa,ll hin,ll hout,int d){//k^d
    sz[x]=1;
    
    for(reg i=0;i<(int)to[x].size();++i){
        int y=to[x][i];
        if(vis[y]||y==fa) continue;
        dfs2(y,x,(hin+mi[d+1]*v[y])%mod,(hout*k+v[y])%mod,d+1);
    }
}
int nowsz;
void divi(int x){
    rt=0;
    cntin=0;cntout=0;
    dfs1(x,0);
    in.clear();out.clear();
    in.ins(0);out.ins(0);
    for(reg i=0;i<(int)to[rt].size();++i){
        int y=to[rt][i];
        if(!vis[y]){
            dfs2(y,rt,v[y],v[y],0);
        }
    }
    
    in.clear();out.clear();
    reverse(to[rt].begin(),to[rt].end());
    for(reg i=0;i<(int)to[rt].size();++i){
        int y=to[rt][i];
        if(!vis[y]){
            dfs2(y,rt,v[y],v[y],0);
        }
    }
    
    vis[rt]=1;
    for(reg i=0;i<(int)to[rt].size();++i){
        int y=to[rt][i];
        if(!vis[y]){
            nowsz=sz[y];
            divi(y);
        }
    }
}
int main(){
    rd(n);rd(mod);rd(k);rd(M);
    for(reg i=1;i<=n;++i) rd(v[i]);
    int x,y;
    mi[0]=1;
    for(reg i=1;i<=n;++i) mi[i]=mi[i-1]*k%mod;
    for(reg i=1;i<=n;++i){
        if(M==0) n0[i]=1,out0[i]=1, 
        else in1[i]=1,out1[i]=1;
    }
    for(reg i=1;i<n;++i){
        rd(x);rd(y);
        to[x].push_back(y);
        to[y].push_back(x);
    }
    nowsz=n;
    divi(1);
}

}
signed main(){
    Miracle::main();
    return 0;
}

/*
   Author: *Miracle*
   Date: 2019/3/12 21:22:28
*/

猜你喜欢

转载自www.cnblogs.com/Miracevin/p/10520071.html