Submission #986818

#TimeUsernameProblemLanguageResultExecution timeMemory
986818vyshniak_nUnique Cities (JOI19_ho_t5)C++17
100 / 100
408 ms76520 KiB
//#pragma optimize("Ofast")
//#pragma optimize("unroll-loops")
//#pragma traget("avx,avx2")

#include <iostream>
#include <cmath>
#include <algorithm>
#include <stdio.h>
#include <cstdint>
#include <cstring>
#include <string>
#include <cstdlib>
#include <vector>
#include <bitset>
#include <map>
#include <queue>
#include <ctime>
#include <stack>
#include <set>
#include <list>
#include <random>
#include <deque>
#include <functional>
#include <iomanip>
#include <sstream>
#include <fstream>
#include <complex>
#include <numeric>
#include <cassert>
#include <array>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <thread>

typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;

#define el '\n'
#define ff first
#define ss second
#define pb push_back
#define pf push_front
#define popb pop_back
#define popf pop_front
#define point pair <ll, ll>
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()

using namespace std;

#include <random>
mt19937 rnd(time(0));

const ll INF = 1e18 + 10;
const ll inf = 1e9 + 10;
const ll N = 4e5 + 10;
const ll mod = 1e9 + 7;
const ll K = 200;
const ll LOG = 8;

vector <ll> gr[N], dist[2];
ll a[N], mx[2][N];

void dfs(ll v, ll pr, ll id, ll h = 1) {
    dist[id][v] = h;
    mx[id][v] = h;
    for (ll to : gr[v]) {
        if (to == pr)
            continue;

        dfs(to, v, id, h + 1);
        mx[id][v] = max(mx[id][v], mx[id][to]);
    }
}

ll cnt[N], cur = 0, ans[N];
vector <ll> st;

void add(ll x) {
    cnt[a[x]]++;
    if (cnt[a[x]] == 1)
        cur++;
}
void del(ll x) {
    cnt[a[x]]--;
    if (cnt[a[x]] == 0)
        cur--;
}

void calc(ll v, ll pr, ll id) {
    vector <point> children;
    for (ll to : gr[v]) {
        if (to == pr)
            continue;

        children.pb({mx[id][to], to});
    }

    sort(rall(children));
    while (children.size() > 1 && !st.empty() && dist[id][st.back()] >= 2 * dist[id][v] - children[1].ff) {
        del(st.back());
        st.popb();
    }

    for (ll i = 0; i < children.size(); i++) {
        ll to = children[i].ss;

        add(v);
        st.pb(v);

        calc(to, v, id);

        while (!st.empty() && dist[id][st.back()] >= 2 * dist[id][v] - children[0].ff) {
            del(st.back());
            st.popb();
        }
    }

    if (dist[id][v] >= dist[id ^ 1][v])
        ans[v] = cur;
}

void solve() {
    ll n, m;
    cin >> n >> m;

    for (ll i = 1; i < n; i++) {
        ll u, v;
        cin >> u >> v;

        gr[u].pb(v);
        gr[v].pb(u);
    }

    for (ll i = 1; i <= n; i++)
        cin >> a[i];

    ll d1, d2;
    dist[0].resize(n + 1);
    dist[1].resize(n + 1);

    dfs(1, 1, 0);
    d1 = max_element(all(dist[0])) - dist[0].begin();
    dfs(d1, d1, 0);
    d2 = max_element(all(dist[0])) - dist[0].begin();
    dfs(d2, d2, 1);

    calc(d1, d1, 0);
    calc(d2, d2, 1);

    for (ll i = 1; i <= n; i++)
        cout << ans[i] << el;
    return;
}
signed main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);

    int tests = 1;
    //cin >> tests;

    while (tests--) 
        solve();
    return 0;
}
/*
*/

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'void calc(ll, ll, ll)':
joi2019_ho_t5.cpp:107:22: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  107 |     for (ll i = 0; i < children.size(); i++) {
      |                    ~~^~~~~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...