Submission #1275489

#	Time	Username	Problem	Language	Result	Execution time	Memory
1275489		MisterReaper	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	410 ms	33408 KiB

joi2019_ho_t5

// File uniquecities.cpp created on 02.10.2025 at 15:34:56
#include <bits/stdc++.h>

using i64 = long long;

#ifdef DEBUG 
    #include "/home/ahmetalp/Desktop/Workplace/debug.h"
#else
    #define debug(...) void(23)
#endif

constexpr int max_N = int(2E5) + 5;

int N, M;
std::vector<int> adj[max_N];
int C[max_N];

int dep[max_N];

void dfs1(int v, int pr) {
    for (auto u : adj[v]) {
        if (u == pr) {
            continue;
        }
        dep[u] = dep[v] + 1;
        dfs1(u, v);
    }
}

int h[max_N];

void dfs2(int v, int pr) {
    h[v] = 0;
    for (auto u : adj[v]) {
        if (u == pr) {
            continue;
        }
        dep[u] = dep[v] + 1;
        dfs2(u, v);
        h[v] = std::max(h[v], h[u] + 1);
    }
}

int now = 0;
int cnt[max_N];

std::vector<int> stk;

void add(int v) {
    stk.emplace_back(v);
    if (cnt[C[v]]++ == 0) {
        now++;
    }
}

void del(int v) {
    if (cnt[C[v]]-- == 1) {
        now--;
    }
}

void roll(int x) {
    while (!stk.empty() && dep[stk.back()] > x) {
        del(stk.back());
        stk.pop_back();
    }
}

int ans[max_N];

void dfs3(int v, int pr) {
    int big = -1, mx1 = -1, mx2 = -1;
    for (auto u : adj[v]) {
        if (u == pr) {
            continue;
        }
        int x = h[u] + 1;
        if (x > mx1) {
            big = u;
            mx2 = mx1;
            mx1 = x;
        } else if (x > mx2) {
            mx2 = x;
        }
    }
    if (big == -1) {
        debug(v, stk);
        ans[v] = std::max(ans[v], now);
        roll(dep[v] - 1);
        return;
    }
    roll(dep[v] - mx2 - 1);
    add(v);
    dfs3(big, v);
    roll(dep[v] - mx1 - 1);
    roll(dep[v] - 1);
    debug(v, stk);
    ans[v] = std::max(ans[v], now);
    add(v);
    for (auto u : adj[v]) {
        if (u == pr || u == big) {
            continue;
        }
        dfs3(u, v);
        roll(dep[v] - 1);
        add(v);
    }
    roll(dep[v] - 1);
}

void solve(int v) {
    dep[v] = 0;
    dfs2(v, -1);
    dfs3(v, -1);
    debug();
}

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

    std::cin >> N >> M;

    for (int i = 1; i < N; ++i) {
        int A, B;
        std::cin >> A >> B;
        --A, --B;
        adj[A].emplace_back(B);
        adj[B].emplace_back(A);
    }

    for (int i = 0; i < N; ++i) {
        std::cin >> C[i];
        --C[i];
    }

    dep[0] = 0;
    dfs1(0, -1);
    int d0 = std::max_element(dep, dep + N) - dep;
    dep[d0] = 0;
    dfs1(d0, -1);
    int d1 = std::max_element(dep, dep + N) - dep;

    solve(d0);
    solve(d1);

    for (int i = 0; i < N; ++i) {
        std::cout << ans[i] << '\n';
    }

    return 0;
}

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