[Gym] I - Sky Tax

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : I - Sky Tax
Source            : Codeforces
Category          : Graph Theory
Algorithm         : LCA
Verdict           : Accepted

#include <bits/stdc++.h>  
  
using namespace std;  
  
#define int                 long long int  
  
static const int maxn = 2e5 + 5;  
static const int logn = 20;  
  
  
vector <int> graph[maxn];  
int depth[maxn], father[maxn][logn];  
int subtreeSize[maxn];  
  
void dfs(int u, int p = -1)  
{  
      for (int i = 1; i < logn; i++) father[u][i] = father[ father[u][i-1] ][i-1];  
      subtreeSize[u] = 1;  
      for (int v : graph[u])  
      {  
            if (v == p) continue;  
            depth[v] = depth[u] + 1;  
            father[v][0] = u;  
            dfs(v, u);  
            subtreeSize[u] += subtreeSize[v];  
      }  
}  
  
int get_kth_parrent(int u, int k)  
{  
      for (int i = logn-1; i >= 0; i--)  
      {  
            if ((1 << i) <= k)  
            {  
                  u = father[u][i];  
                  k -= (1 << i);  
            }  
      }  
      return u;  
}  
  
int lca(int u, int v)  
{  
      if (depth[u] < depth[v]) swap(u, v);  
      for (int i = logn-1; i >= 0; i--)  
      {  
            if (depth[ father[u][i] ] >= depth[v])  
            {  
                  u = father[u][i];  
            }  
      }  
      if (u == v) return u;  
      for (int i = logn-1; i >= 0; i--)  
      {  
            if (father[u][i] != father[v][i])  
            {  
                  u = father[u][i];  
                  v = father[v][i];  
            }  
      }  
      return father[u][0];  
}  
  
  
signed main()  
{  
      ios_base::sync_with_stdio(false);  
      cin.tie(nullptr);  
      cout.tie(nullptr);  
  
      #ifndef ONLINE_JUDGE  
            freopen("in.txt", "r", stdin);  
      #endif // ONLINE_JUDGE  
  
  
      int tc;  
      cin >> tc;  
      for (int tcase = 1; tcase <= tc; tcase++)  
      {  
            int tnode, tquery, root;  
            cin >> tnode >> tquery >> root;  
            for (int e = 1; e < tnode; e++)  
            {  
                  int u, v;  
                  cin >> u >> v;  
                  graph[u].push_back(v);  
                  graph[v].push_back(u);  
            }  
            depth[root] = 1;  
            dfs(root);  
            cout << "Case #" << tcase << ":\n";  
            while (tquery--)  
            {  
                  int type;  
                  cin >> type;  
                  if (type == 0)  
                  {  
                        int nroot;  
                        cin >> nroot;  
                        root = nroot;  
                  }  
                  else  
                  {  
                        int x;  
                        cin >> x;  
                        int lcax = lca(root, x);  
                        if (x == root)  
                        {  
                              cout << tnode << '\n';  
                              continue;  
                        }  
                        if (lcax == root or (lcax != root && lcax != x))  
                        {  
                              cout << subtreeSize[x] << '\n';  
                              continue;  
                        }  
                        assert(lcax == x);  
                        int d = depth[root] + depth[x] - 2*depth[lcax];  
                        int par = get_kth_parrent(root, d - 1);  
                        int ans = tnode - subtreeSize[par];  
                        cout << ans << '\n';  
                  }  
            }  
            for (int i = 0; i < maxn; i++)  
            {  
                  graph[i].clear();  
                  depth[i] = 0;  
                  subtreeSize[i] = 0;  
                  fill(begin(father[i]), end(father[i]), 0);  
            }  
      }  
  
}