[Spoj] ADAAPPLE – Ada and Apple

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : ADAAPPLE – Ada and Apple
Source            : Spoj
Category          : Graph Theory
Algorithm         : Heavy Light Decomposition, Binary Indexed Tree
Verdict           : Accepted

#include "bits/stdc++.h"
 
using namespace std;
 
using pii = pair <int, int>;
 
static const int maxn = 3e5 + 5;
static const int logn = 19;
 
int chainNo, chainHead[maxn], chainInd[maxn], chainPos[maxn], chainSize[maxn],
    baseArray[maxn], posBase[maxn], ptr, depth[maxn], father[maxn][logn],
    subtreeSize[maxn], startNode[maxn], endNode[maxn], S, E, edgeCost[maxn];
 
vector <int> graph[maxn];
int arr[maxn];
 
inline void dfs(int u = 1, 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];
      }
}
 
inline 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];
}
 
inline void hld(int u = 1, int p = -1)
{
      if (chainHead[chainNo] == -1) chainHead[chainNo] = u;
      chainInd[u] = chainNo;
      ptr++;
      baseArray[ptr] = arr[u];
      posBase[u] = ptr;
      int specialChild = -1;
      int maxSize = -1;
      for (int v : graph[u])
      {
            if (v == p) continue;
            if (subtreeSize[v] > maxSize) maxSize = subtreeSize[v], specialChild = v;
      }
      if (specialChild != -1) hld(specialChild, u);
      for (int v : graph[u])
      {
            if (v == p || v == specialChild) continue;
            chainNo++;
            hld(v, u);
      }
}
 
int tree[maxn];
int N;
 
inline void update(int pos, int val)
{
      while (pos <= N)
      {
            tree[pos] += val;
            pos += pos & -pos;
      }
}
 
inline int read(int pos)
{
      int sum = 0;
      while (pos > 0)
      {
            sum += tree[pos];
            pos -= pos & -pos;
      }
      return sum;
}
 
inline int query(int l, int r)
{
      return read(r) - read(l-1);
}
 
inline int queryUp(int u, int v)
{
      // u : always lowest node
      // v : lca
      int uChain, vChain = chainInd[v], sum = 0;
      while (1)
      {
            uChain = chainInd[u];
            if (uChain == vChain)
            {
                  int l = posBase[v];
                  int r = posBase[u];
                  if (l > r) break;
                  sum += query(l, r);
                  break;
            }
            int l = posBase[ chainHead[uChain] ];
            int r = posBase[u];
            sum += query(l, r);
            u = father[ chainHead[uChain] ][0];
      }
      return sum;
}
 
inline pii query_hld(int u, int v)
{
      int lca = LCA(u, v);
      int dist = depth[u] + depth[v] - 2 * depth[lca] + 1;
      int sum = queryUp(u, lca) + queryUp(v, lca) - queryUp(lca, lca);
      return pii(dist, sum);
}
 
inline void init()
{
      ptr = 0;
      chainNo = 0;
      for (int i = 0; i < maxn; i++)
      {
            graph[i].clear();
            subtreeSize[i] = 0;
            chainHead[i] = -1;
            chainInd[i] = 0;
            depth[i] = 0;
            baseArray[i] = 0;
            posBase[i] = 0;
            for (int j = 0; j < logn; j++) father[i][j] = 0;
      }
}
 
inline void changeNodeCost(int node, int pval, int val)
{
      baseArray[ posBase[node] ] = val;
      if (pval == 1) update(posBase[node], -pval);
      update(posBase[node], val);
}
 
int main()
{
      //freopen("in.txt", "r", stdin);
 
      init();
 
      int tnode, tquery;
      cin >> tnode >> tquery;
      string s;
      cin >> s;
      for (int i = 1; i < tnode; i++)
      {
            int u, v;
            cin >> u >> v;
            u += 1, v += 1;
            graph[u].push_back(v);
            graph[v].push_back(u);
      }
      for (int i = 0; i < tnode; i++) arr[i+1] = s[i] - '0';
      depth[1] = 1;
      dfs();
      hld();
      N = ptr;
      for (int i = 1; i <= N; i++) update(posBase[i], baseArray[ posBase[i] ]);
      while (tquery--)
      {
            int type;
            cin >> type;
            if (type == 0) // update
            {
                  int p;
                  cin >> p;
                  p += 1;
                  int pval = arr[p];
                  if (arr[p] == 0) arr[p] = 1;
                  else arr[p] = 0;
                  changeNodeCost(p, pval, arr[p]);
            }
            else
            {
                  int u, v;
                  cin >> u >> v;
                  u += 1, v += 1;
                  pii ans = query_hld(u, v);
                  if (ans.second == 0 || (ans.first == ans.second)) cout << "YES\n";
                  else cout << "NO\n";
            }
      }
}