1824C - LuoTianyi and XOR-Tree - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
#define int long long int
#define ll long long int
#define pb push_back
int mod = 1e9+7;
int INF = 4e18;
// int INF = 1e9;
void print(int x){
    cout << x << "\n";
}
void print(int x, int y){
    cout << x << ' ' << y << "\n";
}
void print(int x, int y, int z){
    cout << x << ' ' << y << ' ' << z << "\n";
}
void print(vector<int> x){
    for (auto val: x) cout << val << ' ';
    cout << endl;
}
void print(array<int, 2> x){
    cout << x[0] << ' ' << x[1] << endl;
}
void print(vector<vector<int>> x){
    for (auto val: x) print(val);
    
}
void print(vector<array<int, 2>> x){
    for (auto [a, b]: x) cout << a << ' ' << b << endl;
 
}
 
template <class T>
void print(const T& x) {
    for (auto val : x) {
        cout << val << " ";
    }
    cout << "\n";
}
 
void print(set<int> x){
    for (auto val: x) cout << val << ' ';
    cout << endl;
}
void print(multiset<int> x){
    for (auto val: x) cout << val << ' ';
    cout << endl;
}
void print(map<int, int> &x){
    for (auto &[a, b]: x) cout << a << ' ' << b << endl;
}
int log2_floor(int i) {
    return i ? __builtin_clzll(1) - __builtin_clzll(i) : -1;
}
int inv(int a, int b){
    return 1<a ? b - inv(b%a,a)*b/a : 1;
}
int inv(int a){
    return inv(a, mod);
}
vector<int> fact;
vector<int> ifact;
int bin(int N, int K){ 
    if (K>N) return 0;
    if (K<0) return 0;
    if (N<0) return 0;
    int res = 1;
    res *= fact[N]; res %= mod;
    res *= ifact[K]; res %= mod;
    res *= ifact[N-K]; res %= mod;
    return res;
}
void inv_init(int C = 2500000){
    fact.resize(C+1); fact[0] = 1;
    ifact.resize(C+1);
    // invn.resize(C+1);
    for (int i=1; i<=C; i++) fact[i] = (i*fact[i-1])%mod;
    ifact[C] = inv(fact[C]);
    for (int i=C-1; i>=0; i--) ifact[i] = ((i+1)*ifact[i+1])%mod;
}
int binpow(int a, int b) {
    a %= mod;
    long long res = 1;
    while (b > 0) {
        if (b & 1)
            res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res;
}
int n;
int res;
vector<int> v;
vector<vector<int>> adj;
vector<int> leafxor;
vector<int> visited;
vector<int> setmap;
vector<set<int>> s;
void dfs(int x){
    visited[x] = 1;
    leafxor[x] ^= v[x];
    vector<int> mind;
    for (int y: adj[x]){
        if (visited[y]) continue;
        leafxor[y] ^= leafxor[x];
        dfs(y);
        mind.pb(setmap[y]);
    }
    int N = mind.size();
    if (mind.size()==0){
        set<int> nset = {leafxor[x]};
        setmap[x] = s.size();
        s.pb(nset);
        // cout << x << ' ' << leafxor[x] << ' ' << setmap[x] << '\n';
        return;
    }
    sort(mind.begin(), mind.end(), [&](int x, int y){
        return s[x].size() > s[y].size();
    });
    int ok = 1;
    set<int> newval;
    for (int i=0; i<N; i++){
        if (i>0){
            for (int x: s[mind[i]]){
                if (s[mind[0]].find(x) != s[mind[0]].end()) ok = 0;
                if (newval.find(x) != newval.end()) ok = 0;
                newval.insert(x);
            }
        }
    }
    if (ok){

        // cout << mind[0] << "\n";
        // exit(0);
        setmap[x] = mind[0];
        
        // cout << mind[0] << "\n";
        for (int y: newval){
            s[mind[0]].insert(y);
        }    
        // cout << x << ' ' << mind[0] << "\n";
        // // for (int y: s[mind[0]]){
        // //     cout << y << ' ';
        // // }
        // cout << '\n';
    }
    else{
        // cout << "here " << x << "\n";
        map<int, int> cnts;
        set<int> adds;
        int maxcnt = 0;
        for (int ind: mind){
            // cout << ind << ' ';
            for (int y: s[ind]){
                cnts[y]++;
                maxcnt = max(maxcnt, cnts[y]);
            }
        }
        // cout << "\n";
        // cout << maxcnt << "\n";
        for (auto [y, xcnt]: cnts){
            if (xcnt!=maxcnt){
                res += xcnt;
            }
            else{
                res += maxcnt-1;
                adds.insert(y);
            }
        }
        res -= maxcnt-1;
        setmap[x] = s.size();
        s.pb(adds);
        // if (x==1){
        //     for (int y: adds){
        //         cout << y << ' ';
        //     }
        //     cout << "\n";
        // }
    }
}
void solve(){
    int n; cin >> n;
    v.assign(n, {});
    adj.assign(n, {});
    visited.assign(n, {});
    leafxor.assign(n, {});
    setmap.assign(n, {});
    res = 0;
    for (int i=0; i<n; i++) cin >> v[i];
    for (int i=0; i<n-1; i++){
        int x, y; cin >> x >> y; x--; y--;
        adj[x].pb(y); adj[y].pb(x);
    }
    dfs(0);
    // cout << setmap[1] << ' ' << setmap[5] << '\n';
    cout << res + s[setmap[0]].size() - (s[setmap[0]].find(0) != s[setmap[0]].end()) << "\n";


}
int32_t main() {    
    ios_base::sync_with_stdio(0);
    cin.tie(nullptr);
    // inv_init();
    int tt = 1; 
    // cin >> tt;
    for (int xx=1; xx<=tt; xx++){
        // cout << "Case " << xx << ": ";
        solve();
    }
    return 0;
    
    
}