Submission #2336232

---

Source Code Expand

#include<bits/stdc++.h>
using namespace std;

#define int long long

#define rep(i,n) for(int i=0;i<(n);i++)
#define pb push_back
#define all(v) (v).begin(),(v).end()
#define fi first
#define se second
typedef vector<int>vint;
typedef pair<int,int>pint;
typedef vector<pint>vpint;

template<typename A,typename B>inline void chmin(A &a,B b){if(a>b)a=b;}
template<typename A,typename B>inline void chmax(A &a,B b){if(a<b)a=b;}

int N;
vint G[222222];
string S;
void dfs(int v,int p,int d,vint &dist){
    dist[v]=d;
    for(auto u:G[v]){
        if(u==p)continue;
        dfs(u,v,d+1,dist);
    }
}


    /*
    vint lim(N);
    rep(i,N)lim[i]=max(distS[i],distT[i])-1;
    rep(i,N)chmin(lim[i],R[i]+1);
*/


int ans;
int R[222222];
int dw[222222];


vint distS,distT;
void dfs2(int v,int p){
    R[v]=0;
    dw[v]=INT_MAX;
    for(auto u:G[v]){
        if(u==p)continue;
        dfs2(u,v);
        chmax(R[v],R[u]+1);
    }
    if(S[v]=='1')dw[v]=0;

    int lim=max(distS[v],distT[v])-1;
    chmin(lim,R[v]+1);
    ans+=max(0ll,lim-dw[v]+1);
    if(p!=-1)chmin(dw[p],max(dw[v]-1,R[v]+1));
}



signed main(){
    cin>>N;
    rep(i,N-1){
        int a,b;
        cin>>a>>b;
        a--;b--;G[a].pb(b);G[b].pb(a);
    }
    cin>>S;

    distS.resize(N);distT.resize(N);
    dfs(0,-1,0,distS);
    int s=0;
    rep(i,N)if(distS[i]>distS[s])s=i;
    dfs(s,-1,0,distS);
    int t=0;
    rep(i,N)if(distS[i]>distS[t])t=i;
    dfs(t,-1,0,distT);

    vint lis;
    rep(i,N){
        if(max(distS[i],distT[i])==(distS[t]+1)/2)lis.pb(i);
    }

    if(lis.size()==1){
        dfs2(lis[0],-1);
    }
    else{
        dfs2(lis[0],lis[1]);
        dfs2(lis[1],lis[0]);
    }
    cout<<ans+1<<endl;


    return 0;
}

Submission Info

Judge Result