Submission #126117

#	Time	Username	Problem	Language	Result	Execution time	Memory
126117		EntityIT	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	1081 ms	54036 KiB

joi2019_ho_t5

#include<bits/stdc++.h>

using namespace std;

#define fi first
#define se second
#define pb push_back

const int N = (int)2e5 + 5, M = N;
int n, m, c[N], diameterTip[2], depth[N], ans[N], cnt[M], curAns, mx[N][2];
pair<int, bool> group[N];
set< pair<int, int> > valid, path;
vector<int> gr[N];

void dfs (int u, int p) {
    depth[u] = depth[p] + 1;
    for (int v : gr[u]) if (v != p) {
        dfs(v, u);
    }
}

void prep (int u, int p) {
    depth[u] = depth[p] + 1;
    mx[u][0] = mx[u][1] = 0;
    for (int v : gr[u]) if (v != p) {
        prep(v, u);
        int tmp = mx[v][0];
        for (int _ = 0; _ < 2; ++_) if (mx[u][_] < tmp) swap(tmp, mx[u][_]);
    }
    mx[u][0] = max(mx[u][0], depth[u]);
}

void add (int u) {
    if (!cnt[ c[u] ]) ++curAns;
    ++cnt[ c[u] ];
}

void subtract (int u) {
    --cnt[ c[u] ];
    if (!cnt[ c[u] ]) --curAns;
}

void calc (int v, int Mx) {
//    cout << "v = " << v << "  Mx = " << Mx << '\n';
    while (path.size() && (*path.begin() ).fi < depth[v] - (mx[v][0] - depth[v]) ) {
        valid.insert(*path.begin() );
        add( (*path.begin() ).se );
        path.erase(path.begin() );
    }
//    for (auto _ : valid) cout << _.se << ' '; cout << '\n';
    while (valid.size() && (*valid.rbegin() ).fi >= depth[v] - Mx) {
        subtract( (*valid.rbegin() ).se );
        valid.erase(--valid.end() );
    }
    while (path.size() && (*path.rbegin() ).fi >= depth[v] - Mx) {
        path.erase(--path.end() );
    }
//    for (auto _ : valid) cout << _.se << ' '; cout << '\n';
}

void solve (int u, int p, int _) {
//    cout << "u = " << u << "  p = " << p << '\n';
//    cout << curAns << "+++\n";

    if (p && depth[u] == mx[u][0]) add(p);
    if (group[u].se == _) ans[u] += curAns;
    if (p && depth[u] == mx[u][0]) subtract(p);

    if (p) path.insert( { depth[p], p } );

    for (int v : gr[u]) if (v != p && mx[v][0] == mx[u][0]) {
        calc(v, max(0, mx[u][1] - depth[u] + 1) ); solve(v, u, _);
        for (int w : gr[u]) if (w != p && w != v) calc(w, max(0, mx[u][0] - depth[u] + 1) ), solve(w, u, _);
        break ;
    }

    if (valid.find( { depth[p], p } ) != valid.end() ) valid.erase( { depth[p], p } ), subtract(p);
    if (path.find( { depth[p], p } ) != path.end() ) path.erase( { depth[p], p } );
}

int main () {
//    freopen("test.INP", "r", stdin);
    scanf("%d %d", &n, &m);
    for (int _ = 1, u, v; _ < n; ++_) {
        scanf("%d %d", &u, &v);
        gr[u].pb(v);
        gr[v].pb(u);
    }
    for (int _ = 1; _ <= n; ++_) scanf("%d", c + _);

    dfs(1, 0);
    for (int u = 1; u <= n; ++u) if (depth[u] > depth[ diameterTip[0] ]) diameterTip[0] = u;
    dfs(diameterTip[0], 0);
    for (int u = 1; u <= n; ++u) group[u] = { depth[u], 0 };
    for (int u = 1; u <= n; ++u) if (depth[u] > depth[ diameterTip[1] ]) diameterTip[1] = u;
    dfs(diameterTip[1], 0);
    for (int u = 1; u <= n; ++u) group[u] = max(group[u], { depth[u], 1 } );

    for (int _ = 0; _ < 2; ++_) {
//        cerr << diameterTip[_] << "==========================\n";
        prep(diameterTip[_], 0);
        solve(diameterTip[_], 0, _);
    }

    for (int u = 1; u <= n; ++u) printf("%d\n", ans[u]);

    return 0;
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:83:10: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
     scanf("%d %d", &n, &m);
     ~~~~~^~~~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:85:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%d %d", &u, &v);
         ~~~~~^~~~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:89:39: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
     for (int _ = 1; _ <= n; ++_) scanf("%d", c + _);
                                  ~~~~~^~~~~~~~~~~~~

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