Submission #973544

#	Time	Username	Problem	Language	Result	Execution time	Memory
973544		efedmrlr	Unique Cities (JOI19_ho_t5)	C++17	4 / 100	1341 ms	274432 KiB

joi2019_ho_t5

// #pragma GCC optimize("O3,Ofast,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include <bits/stdc++.h>

using namespace std;


#define lli long long int
#define MP make_pair
#define pb push_back
#define REP(i,n) for(int i = 0; (i) < (n); (i)++)
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()


void fastio() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
}


const double EPS = 0.00001;
const int INF = 1e9+500;
const int N = 3e5+5;
const int ALPH = 26;
const int LGN = 25;
constexpr int MOD = 1e9+7;
int n = 0, m = 0;
int D1 = 0, D2 = 0;
vector<int> t(N, 0);
vector<vector<int> > adj(N, vector<int>());
vector<vector<int> > mxsub;
vector<int> res(N, 0), dep, mxs;
vector<int> stck;
vector<int> freq;
int ans;
array<int, 2> get_far(int node, int par, int dist = 0) {
    array<int, 2> ret = {dist, node};
    for(auto c : adj[node]) {
        if(c == par) continue;
        ret = max(ret, get_far(c, node, dist + 1));
    }   
    return ret;
}

void pop_el() {
    freq[t[stck.back()]]--;
    if(freq[t[stck.back()]] == 0) ans--;
    stck.pop_back();
}
void add_el(int x) {
    freq[t[x]]++;
    stck.pb(x);
    if(freq[t[x]] == 1) ans++;
}

int prec(int node, int par) {
    array<array<int, 2>, 2> mx;
    mx[0] = {0, 0}; mx[1] = {0, 0};
    for(auto c : adj[node]) {
        if(c == par) continue;
        dep[c] = dep[node] + 1;
        int ret = prec(c, node);
        if(ret > mx[0][0]) {
            swap(mx[0], mx[1]);
            mx[0] = {ret, c};
        }
        else if(ret > mx[1][0]) {
            mx[1] = {ret, c};
        }
    }
    int ind = -1;
    for(int i = 0; i < adj[node].size(); i++) {
        if(adj[node][i] == par) continue;
        if(mx[0][1] == adj[node][i]) {
            mxsub[node][i] = mx[1][0];
            ind = i;
        }
        else {
            mxsub[node][i] = mx[0][0];
        }
    }
    if(ind != -1) {
        swap(mxsub[node][ind], mxsub[node][mxsub[node].size() - 1]);
        swap(adj[node][ind], adj[node][adj[node].size() - 1]);
    }
    mxs[node] = mx[0][0];
    return mx[0][0] + 1;
}

void dfs(int node, int par) {
    // cout << "node:" << node << "\n";
    vector<int> tmp(0);
    while(stck.size() && dep[stck.back()] >= dep[node] - mxs[node]) {
        tmp.pb(stck.back());
        pop_el();
    }
    res[node] = max(res[node], ans);
    // cout << "stck:";
    // for(auto c : stck) cout << c << " ";
    // cout << "\n";
    add_el(node);
    for(int i = 0; i < adj[node].size() - 1; i++) {
        auto c = adj[node][i];
        if(c == par) continue;
        dfs(c, node);
        
    }
    pop_el();
    int last = 0;
    if(mxsub[node].size() > 1 || par == 0) last = mxsub[node].back();

    while(tmp.size() && dep[tmp.back()] < dep[node] - last) {
        add_el(tmp.back());
        tmp.pop_back();
    }
    add_el(node);
    if(mxsub[node].size() > 1 || par == 0) {
        dfs(adj[node].back(), node);
    }
    pop_el();
    while(tmp.size()) {
        add_el(tmp.back());
        tmp.pop_back();
    }

}

void find_res(int rt) {
    mxsub.assign(n + 1, vector<int>());
    dep.assign(n + 1, 0);
    stck.clear();
    mxs.assign(n + 1, 0);
    freq.assign(n + 1, 0);
    ans = 0;
    for(int i = 1; i <= n; i++) {
        mxsub[i].resize(adj[i].size());
    }
    prec(rt, 0);
    dfs(rt, 0);
}

inline void solve() {
    cin>>n>>m;
    REP(i, n - 1) {
        int u, v;
        cin >> u >> v;
        adj[u].pb(v);
        adj[v].pb(u);
    }
    for(int i = 1; i <= n; i++) {
        cin >> t[i];
    }
    D1 = get_far(1, 0)[1];
    D2 = get_far(D1, 0)[1];
    // cout << D1 << " " << D2 << "\n";
    find_res(D1);
    find_res(D2);
    for(int i = 1; i <= n; i++) {
        cout << res[i] << "\n";
    }
}
 
signed main() {

    fastio();
    int test = 1;
    //cin>>test;
    while(test--) {
        solve();
    }
    
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'int prec(int, int)':
joi2019_ho_t5.cpp:73:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   73 |     for(int i = 0; i < adj[node].size(); i++) {
      |                    ~~^~~~~~~~~~~~~~~~~~
joi2019_ho_t5.cpp: In function 'void dfs(int, int)':
joi2019_ho_t5.cpp:103:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  103 |     for(int i = 0; i < adj[node].size() - 1; i++) {
      |                    ~~^~~~~~~~~~~~~~~~~~~~~~

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