Submission #773604

#	Time	Username	Problem	Language	Result	Execution time	Memory
773604		1bin	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	307 ms	38848 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
 
using namespace std;
 
#define all(v) v.begin(), v.end()
typedef long long ll;
const int NMAX = 2e5 + 5;
int n, m, C[NMAX], D[NMAX], a, b, mxd, ix, r1, r2, ans[NMAX], cur, cnt[NMAX], mx[NMAX][2], idx[NMAX][2];
vector<int> adj[NMAX];
 
void dfs2(int x, int p, int d){
    if(d > mxd) mxd = d, ix = x;
    for(int&nx : adj[x])
        if(nx != p) dfs2(nx, x, d + 1);
    return;
}
 
void go(int x, int p){
    mx[x][0] = mx[x][1] = 0;
    for(int& nx : adj[x])
        if(nx != p){
            go(nx, x);
            if(mx[nx][0] + 1 > mx[x][1]) mx[x][1] = mx[nx][0] + 1, idx[x][1] = nx;
            if(mx[x][0] < mx[x][1]){
                swap(mx[x][0], mx[x][1]);
                swap(idx[x][0], idx[x][1]);
            }
        }
    return;
}
 
void f(int x, int add){
    if(add){
        if(++cnt[C[x]] == 1) cur++;
    }
    else{
        if(--cnt[C[x]] == 0) cur--;
    }
    return;
}
 
stack<int> st;
void dfs(int x, int p, int d){
    D[d] = x;
    
    auto calc = [&] (int mxd){
        while(st.size() && st.top() >= mxd){
            f(D[st.top()], 0); st.pop();
        }
    };
    if(mx[x][0]){
        int nx = idx[x][0];
        if(mx[x][1]) calc(d - mx[x][1]);
        st.emplace(d); f(x, 1);
        dfs(nx, x, d + 1);
        
        calc(d - mx[x][0]);
        ans[x] = max(ans[x], cur);
        for(int& nx : adj[x])
            if(nx != p && nx != idx[x][0]){
                if(st.empty() || st.top() != d) st.emplace(d), f(x, 1);
                dfs(nx, x, d + 1);
            }
        if(st.size() && st.top() == d) st.pop(), f(x, 0);
    }
    else ans[x] = max(ans[x], cur);
    return;
}
 
int main(void){
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    
    cin >> n >> m;
    for(int i = 1; i < n; i++){
        cin >> a >> b;
        adj[a].emplace_back(b);
        adj[b].emplace_back(a);
    }    
    for(int i = 1; i <= n; i++) cin >> C[i];
    dfs2(1, -1, 0);
    r1 = ix;
    mxd = 0; dfs2(r1, -1, 0);
    r2 = ix;
    
    go(r1, -1); dfs(r1, -1, 0);
    go(r2, -1); dfs(r2, -1, 0);
    for(int i = 1; i <= n; i++) cout << ans[i] << '\n';
    return 0;
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32