Submission #12138201

---

Source Code Expand

#include<bits/stdc++.h>

int main(){
    using namespace std;
    unsigned long N;
    cin >> N;
    vector<unsigned long> c(N);
    for(auto& i : c){
        cin >> i;
        --i;
    }
    vector edge(N, vector(0, 0UL));
    for(unsigned long i{1}, a, b; i < N; ++i){
        cin >> a >> b;
        edge[--a].push_back(--b);
        edge[b].push_back(a);
    }
    vector euler_tour(N, vector(0, 0UL));
    vector<unsigned long> sz(N);
    [dfs_impl = [&c, &edge, &euler_tour, &sz, &N](auto f, unsigned long now, unsigned long prev, unsigned long& t) -> unsigned long {
        sz[now] = 1;
        euler_tour[c[now]].push_back(1);
        for(const auto& i : edge[now])if(i != prev){
            sz[now] += f(f, i, now, t);
            euler_tour[c[now]].push_back(sz[now]);
        }
        euler_tour[c[now]].push_back(N + 1);
        return sz[now];
    }]{unsigned long t{0};dfs_impl(dfs_impl, 0, 0, t);}();
    for(unsigned long i{0}; i < N; ++i){
        unsigned long cnt{0};
        vector<unsigned long> tmp{0};
        for(const auto& now : euler_tour[i]){
            if(now == 1){
                tmp.push_back(1);
            }else if(now <= N){
                auto t = now - tmp.back();
                cnt += t * (t + 1) / 2;
                tmp.back() = now;
            }else{
                auto t = tmp.back();
                tmp.pop_back();
                tmp.back() += t;
            }
        }
        cnt += (N - tmp.back() + 1) * (N - tmp.back()) / 2;
        cout << N * (N + 1) / 2 - cnt << endl;
    }
    return 0;
}

Submission Info

Judge Result