QTREE6 - Query on a tree VI

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

const int INF = 0x3f3f3f3f;

struct SegTree {
  int N;

  vector<int> tr;

  void init(int n, vector<int> &v) {
    N = n;
    tr.assign(n << 2, 0);
    build(1, 1, n, v);
  }

  inline int join(int a, int b) {
    return a + b;
  }

  void build(int node, int l, int r, vector<int> &v) {
    if (l == r) {
      // tr[node] = 0;
      tr[node] = v[l];
      return;
    }
    int lc = node << 1;
    int mid = l + (r - l) / 2;
    build(lc, l, mid, v);
    build(lc + 1, mid + 1, r, v);
  }

  int query(int node, int l, int r, int idx) {
    if (l == r) return tr[node];
    int lc = node << 1;
    int mid = l + (r - l) / 2;
    if (idx <= mid)
      return tr[node] + query(lc, l, mid, idx);
    else
      return tr[node] + query(lc+1, mid+1, r, idx);
  }

  void update(int node, int l, int r, int ql, int qr, int val) {
    if (r < l or qr < l or r < ql) return;
    if (ql <= l and r <= qr) return void(tr[node] += val);
    int lc = node << 1;
    int mid = l + (r - l) / 2;
    update(lc, l, mid, ql, min(mid, qr), val);
    update(lc+1, mid+1, r, max(mid+1, ql), qr, val);
  }

  void update(int l, int r, int val) {
    return update(1, 1, N, l, r, val);
  }

  int query(int idx) {
    return query(1, 1, N, idx);
  }
};

struct SegTreeMax {
  int N;
  vector<int> tr;

  void init(int n, vector<int> &v) {
    N = n;
    tr.assign(n << 2, 0);
    build(1, 1, n, v);
  }

  inline int join(int a, int b) {
    return max(a, b);
  }

  void build(int node, int l, int r, vector<int> &v) {
    if (l == r) {
      // tr[node] = -INF;
      tr[node] = v[l];
      return;
    }
    int lc = node << 1;
    int mid = l + (r - l) / 2;
    build(lc, l, mid, v);
    build(lc + 1, mid + 1, r, v);
    tr[node] = join(tr[lc], tr[lc+1]);
  }

  void update(int node, int l, int r, int idx, int val) {
    if (l == r) {
      tr[node] = val;
      return;
    }
    int lc = node << 1;
    int mid = l + (r - l) / 2;
    if (idx <= mid)
      update(lc, l, mid, idx, val);
    else
      update(lc+1, mid+1, r, idx, val);
    tr[node] = join(tr[lc], tr[lc+1]);
  }

  int query(int node, int l, int r, int ql, int qr) {
    if (r < l or qr < l or r < ql) return -INF;
    if (ql <= l and r <= qr) return tr[node];
    int lc = node << 1;
    int mid = l + (r - l) / 2;
    return join(query(lc, l, mid, ql, min(mid, qr)),
      query(lc+1, mid+1, r, max(mid+1, ql), qr));
  }

  void update(int idx, int val) {
    return update(1, 1, N, idx, val);
  }

  int query(int l, int r) {
    return query(1, 1, N, l, r);
  }
};

int n;
vector<vector<int>> gr, aux;

int timer;
vector<int> h, sub, pos, dad;

SegTree st[2];
SegTreeMax stMax[2];

vector<int> clr;
vector<int> tr[2], trMax[2];

// st[0] -> armazena os tamanhos da subarvores de cor 0
// st[1] -> armazena os tamanhos da subarvores de cor 1
// stMax[0] -> armazena os índices onde está da cor 0
// stMax[1] -> armazena os índices onde está da cor 1

void dfs(int u, int par = -1) {
  sub[u] = 1;
  for (int &to : gr[u]) if (to != par) {
    dad[to] = u;
    dfs(to, u);
    sub[u] += sub[to];
    if (sub[to] > sub[gr[u][0]] or gr[u][0] == par)
      swap(to, gr[u][0]);
  }
}

void build_hld(int u, int par = -1) {
  pos[u] = ++timer;
  tr[0][pos[u]] = 1;
  tr[1][pos[u]] = sub[u];
  trMax[0][pos[u]] = -INF;
  trMax[1][pos[u]] = pos[u];
  // st[0].update(pos[u], pos[u], 1);
  // st[1].update(pos[u], pos[u], sub[u]);
  // stMax[0].update(pos[u], -INF);
  // stMax[1].update(pos[u], pos[u]);
  for (int to : gr[u]) if (to != par) {
    h[to] = (to == gr[u][0] ? h[u] : to);
    build_hld(to, u);
  }
}

int query(int o) {
  // update_pos -> primeira posição de cor clr[o] ^ 1
  // ou -INF caso não exista
  int u = o, update_pos = -1;
  while (u != -1) {
    update_pos = stMax[clr[o] ^ 1].query(pos[h[u]], pos[u]);
    if (update_pos != -INF)
      break;

    o = h[u];
    u = dad[o];
    // cout << u << ' ' << o << '\n'; // 0 2
    // cout << dad[o] << '\n'; // 0
  }

  if (update_pos == -INF) {
    // caso não exista, faz a query a partir do root
    update_pos = 0;
  } else {
    if (update_pos == pos[u]) {
      // caso update_pos seja o primeiro de uma chain
      update_pos = pos[o];
    } else {
      // caso esteja no meio da chain, pega o filho
      update_pos = update_pos + 1;
    }
  }

  return st[clr[o]].query(update_pos);
}

void update(int o) {
  // diminui da cor antiga
  int sz = st[clr[o]].query(pos[o]);
  int u = dad[o];
  while (u != -1) {
    int update_pos = stMax[clr[o] ^ 1].query(pos[h[u]], pos[u]);
    st[clr[o]].update(update_pos == -INF ? pos[h[u]] : update_pos, pos[u], -sz);
    if (update_pos != -INF) break;
    u = dad[h[u]];
  }

  stMax[clr[o]].update(pos[o], -INF);

  // muda cor
  clr[o] ^= 1;

  stMax[clr[o]].update(pos[o], pos[o]);

  // aumenta da próxima cor
  sz = st[clr[o]].query(pos[o]);
  u = dad[o];
  while (u != -1) {
    int update_pos = stMax[clr[o] ^ 1].query(pos[h[u]], pos[u]);
    st[clr[o]].update(update_pos == -INF ? pos[h[u]] : update_pos, pos[u], sz);
    if (update_pos != -INF) break;
    u = dad[h[u]];
  }
}

int main() {
  cin.tie(0)->sync_with_stdio(0);

  cin >> n;

  gr.assign(n, vector<int>());
  aux.assign(n, vector<int>());
  clr.assign(n, 1);

  for (int i = 0; i < n-1; ++i) {
    int a, b; cin >> a >> b, --a, --b;
    gr[a].push_back(b);
    gr[b].push_back(a);
  }

  sub.assign(n, 0);
  pos.assign(n, 0);
  dad.assign(n, -1);
  h.assign(n, 0);

  tr[0].assign(n + 1, 0);
  tr[1].assign(n + 1, 0);
  trMax[0].assign(n + 1, -INF);
  trMax[1].assign(n + 1, -INF);

  dfs(0);
  build_hld(0);

  // for (int i = 1; i <= n; ++i) {
  //   cout << tr[0][i] << " \n"[i==n];
  // }

  // for (int i = 1; i <= n; ++i) {
  //   cout << tr[1][i] << " \n"[i==n];
  // }

  st[0].init(n, tr[0]);
  st[1].init(n, tr[1]);
  stMax[0].init(n, trMax[0]);
  stMax[1].init(n, trMax[1]);

  // for (int i = 1; i <= n; ++i) {
  //   cout << st[0].query(i) << " \n"[i==n];
  // }

  // for (int i = 1; i <= n; ++i) {
  //   cout << st[1].query(i) << " \n"[i==n];
  // }

  int m; cin >> m;
  while (m--) {
    int t, u; cin >> t >> u, --u;
    if (t == 0) {
      // subir enquanto tiver a mesma cor e depois pegar ans
      cout << query(u) << '\n';
    } else {
      // mudar cor do atual, e atualizar os ancestrais
      update(u);
    }
  }
}
/*
5
1 2
1 3
1 4
1 5
5
0 1
0 2
0 3
0 4
0 5

*/