Submission #1245383

#TimeUsernameProblemLanguageResultExecution timeMemory
1245383damoonUnique Cities (JOI19_ho_t5)C++20
0 / 100
2076 ms71540 KiB
#include <bits/stdc++.h>
using namespace std;

//#pragma GCC optimize("O3,unroll-loops") //main
//#pragma GCC target("avx2") //cf...
//#pragma GCC target("sse4") //Quera

#define ll long long
typedef pair<int,int> pii;
typedef pair<int,pii> pip;
typedef pair<pii,int> ppi;
typedef pair<pii,pii> ppp;
#define f first
#define s second

#define lc 2*id
#define rc 2*id+1
#define all(x) x.begin(),x.end()

#define pb push_back
#define pp pop_back
#define unicorn(x) x.resize(unique(x.begin(),x.end())-x.begin())
#define dig(x,d) ((x&(1ll<<d)) > 0)

string pr(int* vv,int l,int r){for(int i=l;i<r;i++)cout<<vv[i]<<" ";return "";}
string pr( ll* vv,int l,int r){for(int i=l;i<r;i++)cout<<vv[i]<<" ";return "";}
string pr(vector<int> vv){for(auto i:vv)cout<<i<<" ";return "";}
string pr( vector<ll> vv){for(auto i:vv)cout<<i<<" ";return "";}
string pr(pii* vv,int l,int r){for(int i=l;i<r;i++)cout<<"( "<<vv[i].f<<","<<vv[i].s<<" )    ";return "";}
string pr(vector<pii> vv){for(auto i:vv)cout<<"( "<<i.f<<","<<i.s<<" )    ";return "";}
string pr(bool* vv,int l,int r){for(int i=l;i<r;i++)cout<<vv[i];return"";}
ostream& operator <<(ostream& out, pii x){out << "("<<x.f<<", "<<x.s<<") ";return out;}

const int L=4e5+10;
const int inf=1e9+10;
int n,m,d1,d2;
int a[L],par[L],h[L],H[L];
bool sd[L];
vector<int> adj[L];
int last[L],ans[L];

void dfsD(int v){
    H[v] = h[v];
    for(auto u:adj[v]){
        if(u != par[v]){
            h[u] = h[v]+1;
            par[u] = v;
            dfsD(u);
            H[v] = max(H[v],H[u]);
        }
    }
}

void put(int v,int p){
    sd[v] = 1;
    for(auto u:adj[v]){
        if(u != p){
            put(u,v);
        }
    }
}

struct node{
    int lazy,mx,cnt;
    node(){
        lazy = mx = cnt = 0;
    }
};
struct sagi{
    node t[L<<3];
    void spread(int id){
        t[lc].lazy += t[id].lazy;
        t[rc].lazy += t[id].lazy;
        t[id].mx += t[id].lazy;
        t[id].lazy = 0;
    }
    node merge(node a,node b){
        node ans;
        ans.mx = max(a.mx,b.mx);
        ans.cnt += (a.mx == ans.mx ? a.cnt : 0);
        ans.cnt += (b.mx == ans.mx ? b.cnt : 0);
        return ans;
    }
    void build(int id,int l,int r){
        t[id].lazy = t[id].mx = 0;
        t[id].cnt = r-l;
        if(l+1 == r)
            return;
        int mid = (l+r)>>1;
        build(lc,l,mid);
        build(rc,mid,r);
    }

    void update(int id,int l,int r,int l2,int r2,int x){
        if(l2 >= r2)
            return;
        spread(id);
        if(l==l2 and r==r2){
            t[id].lazy += x;
            return;
        }
        int mid = (l+r)>>1;
        if(l2 < mid)
            update(lc,l,mid,l2,min(r2,mid),x);
        if(r2 > mid)
            update(rc,mid,r,max(l2,mid),r2,x);
        spread(lc);
        spread(rc);
        t[id] = merge(t[lc],t[rc]);
    }

    node get(int id,int l,int r,int l2,int r2){
        if(l2 >= r2){
            node ans;
            return ans;
        }
        spread(id);
        if(l==l2 and r==r2)
            return t[id];
        int mid = (l+r)>>1;
        node ans;
        if(l2 < mid)
            ans = merge(ans,get(lc,l,mid,l2,min(r2,mid)));
        if(r2 > mid)
            ans = merge(ans,get(rc,mid,r,max(l2,mid),r2));
        return ans;
    }

    void put(int ind,int x){
        update(1,1,n+1,ind,ind+1,x-get(1,1,n+1,ind,ind+1).mx);
    }

    void prr(){
        for(int i=1;i<=n;i++){
            cout<<get(1,1,n+1,i,i+1).mx<<" ";
        }
        cout<<endl;
    }
}seg;

void dfs(int v,bool S){
    /*
    cout<<v<<endl;
    seg.prr();
    cout<<"---------------"<<endl;
    */
    if(S == sd[v]){
        node d = seg.get(1,1,n+1,1,max(1,2*h[v]-H[v]));
        ans[v] = (d.mx == 1 ? d.cnt : 0);
    }
    if(adj[v].size() == 1 and par[v])
        return;

    pii ind = pii(0,0);
    for(auto u:adj[v]){
        if(u != par[v]){
            ind = max(ind,pii(H[u],u));
        }
    }

    int pre = last[a[v]];
    if(last[a[v]] == inf or seg.get(1,1,n+1,last[a[v]],last[a[v]]+1).mx != 1){
        last[a[v]] = h[v];
        seg.put(h[v],1);
    }
    int mx = h[v];
    for(auto u:adj[v]){
        if(u != par[v] and u != ind.s){
            seg.update(1,1,n+1,max(2*h[v]-H[v],1),h[v],-2);
            dfs(u,S);
            seg.update(1,1,n+1,max(2*h[v]-H[v],1),h[v],+2);
            mx = max(mx,H[u]);
        }
    }

    seg.update(1,1,n+1,max(2*h[v]-mx,1),h[v],-2);
    dfs(ind.s,S);
    seg.update(1,1,n+1,max(2*h[v]-mx,1),h[v],+2);
    if(last[a[v]] != pre)
        seg.put(h[v],0);
    last[a[v]] = pre;
}

int main(){
    //ofstream cout ("out.out");
    //ifstream cin ("in.in");

    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    cin>>n>>m;
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        adj[u].pb(v);
        adj[v].pb(u);
    }
    for(int i=1;i<=n;i++){
        cin>>a[i];
    }

    h[1] = par[1] = 0;
    dfsD(1);
    pii M = pii(0,0);
    for(int i=1;i<=n;i++){
        M = max(M,pii(h[i],i));
    }
    d1 = M.s;

    h[d1] = par[d1] = 0;
    dfsD(d1);
    M = pii(0,0);
    for(int i=1;i<=n;i++){
        M = max(M,pii(h[i],i));
    }
    d2 = M.s;

    int cur=d2, pre=0;
    while(h[cur]*2 >= h[d2]){
        for(auto v:adj[cur]){
            if(v != pre and v != par[cur]){
                put(v,cur);
            }
        }
        sd[cur] = 1;
        cur = par[cur];
    }

    /*
    cout<<"d1,d2: "<<d1<<" "<<d2<<endl;
    cout<<"sd: "<<pr(sd,1,n+1)<<endl;
    */
    
    h[d1] = 1;
    par[d1] = 0;
    fill(last+1,last+n+1,inf);
    seg.build(1,1,n+1);
    dfsD(d1);
    dfs(d1,1);

    h[d2] = 1;
    par[d2] = 0;
    fill(last+1,last+n+1,inf);
    seg.build(1,1,n+1);
    dfsD(d2);
    dfs(d2,0);

    for(int i=1;i<=n;i++){
        cout<<ans[i]<<endl;
    }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...