Submission #1180733

#TimeUsernameProblemLanguageResultExecution timeMemory
1180733patgraUnique Cities (JOI19_ho_t5)C++20
100 / 100
326 ms45276 KiB
#include <bits/stdc++.h>

#define rep(a,b,c) for(auto a = (b); a != (c); a++)
#define repD(a,b,c) for(auto a = (b); a != (c); a--)
#define repIn(a, b) for(auto& a : (b))
#define repIn2(a, b, c) for(auto& [a, b] : (c))

constexpr bool dbg = 0;
#define DEBUG if constexpr(dbg)
#define DC DEBUG std::cerr
#define eol std::endl

#define ll long long
#define pb push_back

using namespace std;

int n, m, tb;
vector<vector<int>> g;
vector<int> c, t, ans, ans2, depth, maxD, atDepth, cntUniques;

void tAdd(int l, int r, int x) {
    if(l > r) return;
    DC << "  tAdd(" << l << ' ' << r << ' ' << x << ")" << eol;
    l += tb; r += tb;
    t[l] += x;
    if(r != l) t[r] += x;
    while(l / 2 != r / 2) {
        if(l % 2 == 0) t[l + 1] += x;
        if(r % 2 == 1) t[r - 1] += x;
        l /= 2; r /= 2;
    }
}

int tQ(int i) {
    DC << "  tQ(" << i << ") = ";
    i += tb;
    int ret = 0;
    while(i) ret += t[i], i /= 2;
    DC << ret << eol;
    return ret;
}

pair<int, int> findDeepest(int v, int p) {
    depth[v] = depth[p] + 1;
    maxD[v] = depth[v];
    pair<int, int> ret = {depth[v], v};
    repIn(u, g[v]) if(u != p) ret = max(ret, findDeepest(u, v)), maxD[v] = max(maxD[v], maxD[u]);
    return ret;
}

void dfsAns(int v, int p) {
    DC << " Entering " << v << eol;
    atDepth[depth[v]] = c[v];
    ans[v] = ans[p];

    auto myReach = max(1, depth[p] - (maxD[v] - depth[p]));
    tAdd(myReach, depth[p] - 1, -1);
    repIn(u, g[v]) if(u != p) {
        auto reach = max(1, depth[v] - (maxD[u] - depth[v]));
        tAdd(reach, depth[p], 1);
    }
    rep(i, myReach, min(myReach + 2, max(myReach, depth[v]))) if(tQ(i) == 0) {
        auto who = atDepth[i];
        cntUniques[who]++;
        if(cntUniques[who] == 1) ans[v]++;
    }

    repIn(u, g[v]) if(u != p) {
        dfsAns(u, v);
    }

    DC << " Leaving " << v << eol;
    rep(i, myReach, min(myReach + 2, max(myReach, depth[v]))) if(tQ(i) == 0) {
        auto who = atDepth[i];
        cntUniques[who]--;
    }
    repIn(u, g[v]) if(u != p) {
        auto reach = max(1, depth[v] - (maxD[u] - depth[v]));
        tAdd(reach, depth[p], -1);
    }
    tAdd(myReach, depth[p] - 1, 1);
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
    cin >> n >> m;
    g.resize(n + 1);
    c.resize(n + 1);
    tb = 1 << (32 - __builtin_clz(n));
    t.resize(2 * tb);
    ans.resize(n + 1); ans2.resize(n + 1); depth.resize(n + 1); atDepth.resize(n + 1), maxD.resize(n + 1);
    cntUniques.resize(m + 1);
    rep(i, 1, n) {
        int a, b;
        cin >> a >> b;
        g[a].pb(b);
        g[b].pb(a);
    }
    rep(i, 1, n + 1) cin >> c[i];
    auto [_, A] = findDeepest(1, 0);
    auto [__, B] = findDeepest(A, 0);
    DC << "Root " << A << eol;
    dfsAns(A, 0);
    swap(ans, ans2);
    DC << "Root " << B << eol;
    findDeepest(B, 0);
    dfsAns(B, 0);

    rep(i, 1, n + 1) cout << max(ans[i], ans2[i]) << '\n';
}

#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...