Submission #99280

#	Time	Username	Problem	Language	Result	Execution time	Memory
99280		imeimi2000	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	777 ms	30644 KiB

joi2019_ho_t5

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;

int n, m;
vector<int> edge[200001];
int col[200001];
int dist[200001];
int dep[200001];

void dfs1(int x, int p) {
    dist[x] = dist[p] + 1;
    for (int i : edge[x]) {
        if (i == p) continue;
        dfs1(i, x);
    }
}

int ans[200001];

void dfs2(int x, int p) {
    dist[x] = dist[p] + 1;
    dep[x] = 0;
    for (int i : edge[x]) {
        if (i == p) continue;
        dfs2(i, x);
        dep[x] = max(dep[x], dep[i] + 1);
    }
}

vector<int> st;
int cnt[200001];
int sum;

void add(int x, int v) {
    x = col[x];
    if (cnt[x] == 0) ++sum;
    cnt[x] += v;
    if (cnt[x] == 0) --sum;
}

void dfs3(int x, int p) {
    int mx1 = -1, mx2 = -1;
    for (int i : edge[x]) {
        if (i == p) continue;
        if (mx1 < dep[i]) mx2 = mx1, mx1 = dep[i];
        else if (mx2 < dep[i]) mx2 = dep[i];
    }
    int mxi = 0;
    for (int i : edge[x]) {
        if (i == p) continue;
        if (dep[i] == mx1) {
            mxi = i;
            break;
        }
    }
    ++mx1; ++mx2;
    while (!st.empty() && dist[x] - dist[st.back()] <= mx2) {
        add(st.back(), -1);
        st.pop_back();
    }
    if (mxi) {
        add(x, 1);
        st.push_back(x);
        dfs3(mxi, x);
    }
    while (!st.empty() && dist[x] - dist[st.back()] <= mx1) {
        add(st.back(), -1);
        st.pop_back();
    }
    ans[x] = max(ans[x], sum);
    for (int i : edge[x]) {
        if (i == p || i == mxi) continue;
        if (st.empty() || st.back() != x) {
            add(x, 1);
            st.push_back(x);
        }
        dfs3(i, x);
        if (!st.empty() && st.back() == x) {
            add(x, -1);
            st.pop_back();
        }
    }
}

void solve(int R) {
    dfs2(R, 0);
    st.clear();
    sum = 0;
    for (int i = 1; i <= m; ++i) cnt[i] = 0;
    dfs3(R, 0);
}

int main() {
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n >> m;
    for (int i = 1; i < n; ++i) {
        int a, b;
        cin >> a >> b;
        edge[a].push_back(b);
        edge[b].push_back(a);
    }
    for (int i = 1; i <= n; ++i) cin >> col[i];
    dfs1(1, 0);
    int P = max_element(dist + 1, dist + (n + 1)) - dist;
    dfs1(P, 0);
    int Q = max_element(dist + 1, dist + (n + 1)) - dist;
    solve(P); solve(Q);
    for (int i = 1; i <= n; ++i) printf("%d\n", ans[i]);
    return 0;
}

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