1797D - Li Hua and Tree - CodeForces Solution

#include "bits/stdc++.h"
using namespace std;

#pragma GCC optimize("Ofast,unroll-loops")
#pragma GCC target("avx,avx2,fma")

#define int long long
#define all(v) (v).begin(), (v).end()
#define sz(v) (int)v.size()

#ifndef ONLINE_JUDGE 
#include "debug.hpp"
#else 
#define debug(...) 114
#endif

void solve() {
      int n, m;
      cin >> n >> m;
      vector<int> imp(n + 1);
      for(int i = 1; i <= n; i++) cin >> imp[i];

      vector<int> adj[n + 1];
      for(int i = 1; i < n; i++) {
            int u, v;
            cin >> u >> v;
            adj[u].push_back(v);
            adj[v].push_back(u);
      }

      vector<int> sub(n + 1), val(n + 1), p(n + 1);
      vector<set<pair<int,int>>> son(n + 1);
      function<void(int, int)> dfs = [&](int u, int par) {
            sub[u] = 1;
            val[u] = imp[u];
            for(auto x : adj[u]) if(x != par) {
                  dfs(x, u);
                  p[x] = u;
                  sub[u] += sub[x];
                  val[u] += val[x];
                  son[u].insert({-sub[x], x});
            }
      };
      dfs(1, 0);

      for(int q = 0; q < m; q++) {
            int type, x;
            cin >> type >> x;
            if(type == 1) {
                  cout << val[x] << '\n';
            } 
            else {
                  if(sub[x] == 1) continue;
                  int sonx = (*son[x].begin()).second;
                  int fax = p[x];
                  son[fax].erase({-sub[x], x});
                  son[x].erase({-sub[sonx], sonx});

                  sub[x] -= sub[sonx];
                  sub[sonx] += sub[x];
                  val[x] -= val[sonx];
                  val[sonx] += val[x];

                  son[fax].insert({-sub[sonx], sonx});
                  son[sonx].insert({-sub[x], x});

                  p[sonx] = p[x];
                  p[x] = sonx;
            }
      }
}

signed main() {
      ios_base::sync_with_stdio(false); 
      cin.tie(NULL);

      int T = 1;
      // cin >> T;
        
      for(int t = 1; t <= T; t++) {
            solve();
      }
}