Submission #1225528

#	Time	Username	Problem	Language	Result	Execution time	Memory
1225528		minhpk	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	190 ms	55040 KiB

joi2019_ho_t5

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

int a, b;
vector<int> z[1000005];
int low[1000005], high[1000005], col[1000005], res[1000005];
vector<int> st;
int u = 0, v = 0, cur;
int cnt[1000005], color = 0;

void pop_node() {
    int node = st.back();
    st.pop_back();
    if (--cnt[col[node]] == 0) --color;
}

void add_node(int node) {
    if (cnt[col[node]]++ == 0) ++color;
    st.push_back(node);
}

void dfs(int i, int par) {
    low[i] = 1;
    for (int p : z[i]) if (p != par) {
        high[p] = high[i] + 1;
        dfs(p, i);
        low[i] = max(low[i], low[p] + 1);
    }
}

bool cmp_low(int A, int B) {
    return low[A] > low[B];
}

void dfs1(int i, int par) {
    if (i != cur && z[i].size() == 1) {
        res[i] = max(res[i], color);
        return;
    }

    int sta = (par != 0 ? 1 : 0);

    while (z[i].size() > sta + 1 && !st.empty() && high[i] - high[st.back()] <= low[z[i][sta+1]]) {
        pop_node();
    }

    add_node(i);
    dfs1(z[i][sta], i);

    while (!st.empty() && high[i] - high[st.back()] <= low[z[i][sta]]) {
        pop_node();
    }

    res[i] = max(res[i], color);

    for (int p : z[i]) {
        if (p == par || p == z[i][sta]) continue;
        while (!st.empty() && high[st.back()] >= high[i]) {
            pop_node();
        }
        add_node(i);
        dfs1(p, i);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> a >> b;
    for (int i = 1; i < a; i++) {
        int x, y;
        cin >> x >> y;
        z[x].push_back(y);
        z[y].push_back(x);
    }
    for (int i = 1; i <= a; i++) {
        cin >> col[i];
    }

    dfs(1, 0);
    for (int i = 1; i <= a; i++)
        if (high[i] > high[u]) u = i;

    memset(high, 0, sizeof high);
    dfs(u, 0);
    for (int i = 1; i <= a; i++)
        if (high[i] > high[v]) v = i;

    cur = u;
    memset(high, 0, sizeof high);
    dfs(u, 0);
    for (int i = 1; i <= a; i++)
        sort(z[i].begin(), z[i].end(), cmp_low);
    fill(cnt, cnt + b + 1, 0);
    color = 0;
    st.clear();
    dfs1(u, 0);

    cur = v;
    memset(high, 0, sizeof high);
    dfs(v, 0);
    for (int i = 1; i <= a; i++)
        sort(z[i].begin(), z[i].end(), cmp_low);
    fill(cnt, cnt + b + 1, 0);
    color = 0;
    st.clear();
    dfs1(v, 0);

    for (int i = 1; i <= a; i++) {
        if (b == 1)
            cout << (res[i] ? 1 : 0) << "\n";
        else
            cout << res[i] << "\n";
    }
    return 0;
}

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