Submission #1180808

#	Time	Username	Problem	Language	Result	Execution time	Memory
1180808		anteknne	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	605 ms	43884 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
typedef long double ld;

#define pb push_back
#define pii pair<int, int>
#define pll pair<ll, ll>
#define st first
#define nd second
#define ordered_set tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>

const int MAXN = 200 * 1000 + 17;
int seg[4 * MAXN];
vector<int> graf[MAXN];
int kolor[MAXN];
int odl[MAXN];
int maxodl[MAXN];
int wyn[MAXN];
int wyn2[MAXN];
int conaodl[MAXN];
unordered_map<int, int> jest;
int n;

void aktualizuj(int p, int k, int a, int b, ll w, int i) {
    if (p > b || k < a) {
        return;
    }
    if (a <= p && k <= b) {
        seg[i] += w;
        return;
    }
    int sr = (p + k)/2;
    aktualizuj(p, sr, a, b, w, i * 2);
    aktualizuj(sr + 1, k, a, b, w, i * 2 + 1);
}

int zapytanie(int p, int k, int ind, int i) {
    if (p == k && k == ind) {
        return seg[i];
    }
    int sr = (p + k)/2;
    ll a = 0;
    if (ind <= sr) {
        a = zapytanie(p, sr, ind, i * 2);
    } else {
        a = zapytanie(sr + 1, k, ind, i * 2 + 1);
    }
    return (a + seg[i]);
}

void dfsodl(int v, int p) {
    maxodl[v] = odl[v];
    for (auto x : graf[v]) {
        if (x != p) {
            odl[x] = odl[v] + 1;
            dfsodl(x, v);
            maxodl[v] = max(maxodl[v], maxodl[x]);
        }
    }
    //cout << v << " " << maxodl[v] << " " << odl[v] << "\n";
}

void akt2 (int zas, int v, int p, int w) {
    aktualizuj(0, n - 1, zas, odl[p] - 1, w * (-1), 1);
    for (auto u : graf[v]) {
        if (u != p) {
            aktualizuj(0, n - 1, max(1, odl[v] - (maxodl[u] - odl[v])), odl[p], w, 1);
        }
    }
}

void DFS (int v, int p) {
    //cout << v << " " << p << "\n";
    conaodl[odl[v]] = kolor[v];
    wyn[v] = wyn[p];
    int zas = max(odl[p] - (maxodl[v] - odl[p]), 1);
    akt2(zas, v, p, 1);

    for (int i = zas; i < min(zas + 2, max(zas, odl[v])); i ++) {
        if (zapytanie(0, n - 1, i, 1) == 0) {
            if (jest[conaodl[i]] == 0) {
                wyn[v] ++;
            }
            jest[conaodl[i]] ++;
        }
    }

    for (auto u : graf[v]) {
        if (u != p) {
             DFS(u, v);
        }
    }

    for (int i = zas; i < min(zas + 2, max(zas, odl[v])); i ++) {
        if (zapytanie(0, n - 1, i, 1) == 0) {
            jest[conaodl[i]] --;
        }
    }

    akt2(zas, v, p, -1);
}

int main(){
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);

    int k, a, b;
    cin >> n >> k;
    for (int i = 0; i < n - 1; i ++){
        cin >> a >> b;
        graf[a].pb(b);
        graf[b].pb(a);
    }
    for (int i = 1; i <= n; i ++){
        cin >> kolor[i];
    }

    odl[1] = 1;
    dfsodl(1, 0);
    int A = 1;
    for (int i = 1; i <= n; i ++){
        if (odl[i] > odl[A])
            A = i;
    }

    odl[A] = 1;
    dfsodl(A, 0);
    int B = 1;
    for (int i = 1; i <= n; i ++){
        if (odl[i] > odl[B])
            B = i;
    }

    DFS(A, 0);
    //cout << "\n";
    for (int i = 1; i <= n; i++){
        wyn2[i] = wyn[i];
    }
    odl[B] = 1;
    dfsodl(B, 0);

    DFS(B, 0);

    //cout << "\n";
    //cout << A << " " << B << "\n";
    for (int i = 1; i <= n; i ++){
        cout << max(wyn[i], wyn2[i]) << "\n";
    }
    return 0;
}

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