Submission #1160219

#TimeUsernameProblemLanguageResultExecution timeMemory
1160219Shadow1Cat Exercise (JOI23_ho_t4)C++20
100 / 100
199 ms72456 KiB
// Programmer: Shadow1

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

using ll = long long;
using ld = long double;
using ull = unsigned long long;
using str = string; // yay python!

#define i64 int64_t
#define show(x) cerr << (#x) << " = " << (x) << '\n';
#define output_vector(v) for(auto &x : v){cout << x << ' ';}cout << '\n';
#define output_pairvector(v) for(auto &x : v){cout << x.first << " " << x.second << '\n';}
#define read_vector(v) for(auto &x : v){cin >> x;}
#define vt vector
#define pq priority_queue
#define pb push_back
#define eb emplace_back
#define pii pair<int,int>
#define fir first
#define sec second
#define sz(x) ll(x.size())
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define int long long
#define discretize(x) sort(x.begin(), x.end()); x.erase(unique(x.begin(), x.end()), x.end());
// 
// 

const int maxn = 2e5 + 5;
vector<int> depth(maxn), a(maxn), adj[maxn];
int twok[maxn][25]; //preset to -1

vector<int> p, ranks, set_size;

const int N = maxn;
void build() {
    p.assign(N+2, 0);
    ranks.assign(N+2, 0);
    set_size.assign(N+2, 1);
    for(int i=1; i<=N; ++i)
        p[i] = i;
}

int find_set(int i) {
    if(p[i] == i)
        return i;
    return p[i] = find_set(p[i]);
}

bool same(int i, int j) {
    return find_set(i) == find_set(j);
}

void unite(int i, int j) {
    if(same(i,j)) return;
    int x = find_set(i), y = find_set(j);
    if(ranks[x] > y)
            swap(x,y);
    p[x] = y;
    if(ranks[x] == ranks[y])
            ++ranks[y];
    set_size[y] += set_size[x];
}

int find_size(int x) {
    return set_size[find_set(x)];
}

void dfs(int x, int p) {
    for (int k = 0; k < 25; ++k) {
        if (twok[x][k] == -1) break;
        twok[x][k+1] = twok[twok[x][k]][k];
    }
    for (int it:adj[x]) {
        if(it == p) continue;
        twok[it][0] = x;
        depth[it] = depth[x] + 1;
        dfs(it,x);
    }
}
int kpar(int x, int k){
    for(int j=0; j<25; j++){
        if(x==-1) break;
        if(k&(1<<j)) x=twok[x][j];
    }
    return x;
}


int lca(int a,int b){
    if (depth[a] < depth[b]) swap(a,b);
    a = kpar(a, depth[a] - depth[b]);
    if (a == b) return a; // edge case where b is an ancestor of a
    for (int k=24;k>=0;k--){
        if (twok[a][k] != twok[b][k]){
            a = twok[a][k]; b = twok[b][k];
        }
    }
    return twok[a][0];
}

int dist(int a, int b) {
    return depth[a] + depth[b] - 2*depth[lca(a, b)];
}

void solve() {
    int n;
    cin >> n;
    memset(twok, -1, sizeof(twok));
    for(int i=1; i<=n; ++i)
        cin >> a[i];
    for(int i=0; i<n-1; ++i) {
        int u, v;
        cin >> u >> v;
        adj[a[u]].push_back(a[v]);
        adj[a[v]].push_back(a[u]);
    }
    build();
    dfs(1, 0);
    vector<int> dp(n+1);
    for(int i=1; i<=n; ++i) {
        for(auto &x : adj[i]) {
            x = find_set(x);
            if(x < i) {
                p[x] = i;
                dp[i] = max(dp[i], dp[x] + dist(i, x));
            }
        }
    }
    cout << dp[n] << '\n';
}






// CHECK YOUR OVERFLOWS!!!! 
signed main() {
    // freopen("output.txt", "w", stdout);
    // freopen("input.txt", "r", stdin);
    ios::sync_with_stdio(false);
    cin.tie(NULL); 
    int T = 1;
    // cin >> T;
    for(int tc = 1; tc <= T; ++tc) {
        // cout << "Case #" << tc << ": ";
        solve();
    }

    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...