제출 #1321847

#제출 시각아이디문제언어결과실행 시간메모리
1321847ivan_alexeevUnique Cities (JOI19_ho_t5)C++17
32 / 100
991 ms70492 KiB
#include <bits/stdc++.h>

using namespace std;

#ifndef lisie_bimbi
#define endl '\n'
#pragma GCC optimize("O3")
#pragma GCC target("avx,avx2,bmi2,fma")
#endif

using ll = long long;
const ll inf = 1'000'000'000;

const int N = 200'000;

struct node{
    int mn;
    int cnt;
    int mod;
};

node merge(node a, node b){
    node ans;
    ans.mod = 0;
    if(a.mn < b.mn){
        ans.mn = a.mn;
        ans.cnt = a.cnt;
    }else if(b.mn < a.mn){
        ans.mn = b.mn;
        ans.cnt = b.cnt;
    }else{
        ans.mn = a.mn;
        ans.cnt = a.cnt + b.cnt;
    }
    return ans;
}

struct segtree{
    int n;
    node d[4 * N];
    void push(int u, int l, int r){
        if(d[u].mod != 0){
            if(l + 1 != r){
                d[u * 2 + 1].mod += d[u].mod;
                d[u * 2 + 2].mod += d[u].mod;
            }
            d[u].mn += d[u].mod;
            d[u].mod = 0;
        }
    }
    void build(int u, int l, int r){
        d[u].mn = 0;
        d[u].cnt = r - l;
        d[u].mod = 0;
        if(l + 1 == r){
            
        }else{
            int m = (l + r) / 2;
            build(u * 2 + 1, l ,m);
            build(u * 2 + 2, m, r);
        }
    }
    int ql, qr, dd;
    void update(int u, int l, int r){
        if((ql >= r) || (qr <= l)){
            return;
        }
        if((ql <= l) && (r <= qr)){
            d[u].mod += dd;
            return;
        }
        push(u, l, r);
        int m = (l + r) /  2;
        update(u * 2 + 1, l, m);
        update(u * 2 + 2, m, r);
        push(u * 2 + 1, l, m);
        push(u * 2 + 2, m, r);
        d[u] = merge(d[u * 2 + 1], d[u * 2 + 2]);
    }
    node get(int u, int l, int r){
        if((ql >= r) || (qr <= l)){
            return {inf, 0};
        }
        push(u, l, r);
        if((ql <= l) && (r <= qr)){
            return d[u];
        }
        int m = (l + r) / 2;
        return merge(get(u * 2 + 1, l, m), get(u * 2 + 2, m, r));
    }
    int gg(int l, int r){
        ql = l;
        qr = r;
        auto [mn, cnt, mod] = get(0, 0, n);
        if(mn == 0){
            return cnt;
        }else{
            return 0;
        }
    }
    void upd(int l, int r, int dd2){
        ql = l;
        qr = r;
        dd = dd2;
        update(0, 0, n);
    }
    void vivo(int u, int l, int r){
        if(l + 1 == r){
            cout << d[l].mn << ' ';
        }else{
            int m = (l + r) / 2;
            vivo(u * 2 + 1, l, m);
            vivo(u * 2 + 2, m, r);
        }
        if(u == 0){
            cout << endl;
        }
    }
};

vector<vector<int>> v;
vector<int> c;

pair<int, int> furthest(int u, int par){
    int mx = 0, z = u;
    for(auto i : v[u]){
        if(i != par){
            auto [mx2, z2] = furthest(i, u);
            if(mx2 + 1 > mx){
                mx = mx2 + 1;
                z = z2;
            }
        }
    }
    return {mx, z};
}

void calcdist(int u, int par, vector<int> &d){
    if(par == -1){
        d[u] = 0;
    }else{
        d[u] = d[par] + 1;
    }
    for(auto i : v[u]){
        if(i != par){
            calcdist(i, u, d);
        }
    }
}

void calcleaf(int u, int par, vector<int> &d){
    for(auto i : v[u]){
        if(i != par){
            calcleaf(i, u, d);
            d[u] = max(d[u], d[i] + 1);
        }
    }
}

int n;
segtree st;
set<int> distinct;

void calcans(int u, int par, int level, vector<int> &d, vector<int> &ans){
    st.upd(max(level - d[u], 0), level + 1, 1);
    ans[u] = st.gg(0, level);
    st.upd(max(level - d[u], 0), level + 1, -1);
    bool ins = 0;
    if(distinct.count(c[u])){
        ins = 1;
        st.upd(level, level + 1, 1);
    }else{
        distinct.insert(c[u]);
    }
    int mx1 = 0, mx2 = 0;
    for(auto i : v[u]){
        if(i != par){
            if(d[i] + 1 > mx1){
                mx2 = mx1;
                mx1 = d[i] + 1;
            }else if(d[i] + 1 > mx2){
                mx2 = d[i] + 1;
            }
        }
    }
    for(auto i : v[u]){
        if(i != par){
            if(d[i] + 1 != mx1){
                st.upd(max(0, level - mx1), level, 1);
                calcans(i, u, level + 1, d, ans);
                st.upd(max(0, level - mx1), level, -1);
            }else{
                st.upd(max(0, level - mx2), level, 1);
                calcans(i, u, level + 1, d, ans);
                st.upd(max(0, level - mx2), level, -1);
            }
        }
    }
    if(ins){
        st.upd(level, level + 1, 1);
    }else{
        distinct.erase(c[u]);
    }
}

vector<int> calc(int s){
    vector<int> a(n);
    calcleaf(s, -1, a);
    vector<int> ans(n);
    st.n = n;
    st.build(0, 0, n);
    calcans(s, -1, 0, a, ans);
    return ans;
}

void solve(){
    int m;
    cin >> n >> m;
    v.resize(n);
    for(int i = 0; i < n - 1; i++){
        int x, y;
        cin >> x >> y;
        x--;y--;
        v[x].push_back(y);
        v[y].push_back(x);
    }
    c.resize(n);
    for(int i = 0; i < n; i++){
        cin >> c[i];
    }
    auto [_, s1] = furthest(0, -1);
    auto [__, s2] = furthest(s1, -1);
    vector<int> d1(n), d2(n);
    calcdist(s1, -1, d1);
    calcdist(s2, -1, d2);
    vector<int> ans1, ans2;
    ans1 = calc(s1);
    ans2 = calc(s2);
    for(int i = 0; i < n; i++){
        if(d1[i] > d2[i]){
            cout << ans1[i] << endl;
        }else{
            cout << ans2[i] << endl;
        }
    }
}

signed main(){
#ifdef lisie_bimbi
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
#else
#endif
    cin.tie(nullptr)->sync_with_stdio(false);
    int t = 1;
    //cin >> t;
    while(t--){
        solve();
    }
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...