Submission #591344

#	Time	Username	Problem	Language	Result	Execution time	Memory
591344		piOOE	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	349 ms	38424 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;

const int N = 200001;

vector<int> g[N], stk;

int c[N], cnt[N], cnt_unique = 0, ans[N], depth[N], dp[N], diam;

void pop() {
    int x = stk.back();
    stk.pop_back();
    if (--cnt[c[x]] == 0) {
        cnt_unique -= 1;
    }
}

void push(int x) {
    stk.push_back(x);
    if (++cnt[c[x]] == 1) {
        cnt_unique += 1;
    }
}

void dfs1(int v, int p) {
    dp[v] = depth[v] = depth[p] + 1;
    if (dp[v] > dp[diam]) {
        diam = v;
    }
    for (int to: g[v]) {
        if (to != p) {
            dfs1(to, v);
            dp[v] = max(dp[v], dp[to]);
        }
    }
}

void dfs(int v, int p) {
    if (p != 0) {
        push(p);
    }
    for (int t = 0; t < 2; ++t) {
        for (int i = t; i < g[v].size(); ++i) {
            if (g[v][i] != p && (g[v][t] == p || dp[g[v][i]] > dp[g[v][t]])) {
                swap(g[v][t], g[v][i]);
            }
        }
    }
    int mx1 = g[v][0], mx2 = g[v].size() == 1 || g[v][1] == p ? 0 : g[v][1];
    for (int to: g[v]) {
        if (to != p) {
            int len = (to == mx1 ? dp[mx2] : dp[mx1]) - depth[v];
            while (!stk.empty() && depth[v] - depth[stk.back()] <= len) {
                pop();
            }
            dfs(to, v);
        }
    }
    while (!stk.empty() && dp[v] - depth[v] >= depth[v] - depth[stk.back()]) {
        pop();
    }
    ans[v] = max(ans[v], cnt_unique);
}


int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, m;
    cin >> n >> m;
    for (int i = 1; i < n; ++i) {
        int x, y;
        cin >> x >> y;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    for (int i = 1; i <= n; ++i) {
        cin >> c[i];
    }
    dfs1(diam = 1, 0);
    for (int t = 0, x; t < 2; ++t) {
        //t == 0 ? go from one end : from another
        dfs1(x = diam, 0);
        dfs(x, 0);
    }
    for (int i = 1; i <= n; ++i) {
        cout << ans[i] << "\n";
    }
    return 0;
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'void dfs(int, int)':
joi2019_ho_t5.cpp:44:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   44 |         for (int i = t; i < g[v].size(); ++i) {
      |                         ~~^~~~~~~~~~~~~

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