USACO 2019 December Contest, Gold Problem 2. Milk Visits - STIVA + LCA

#include <bits/stdc++.h>

using namespace std;

ifstream fin("milkvisits.in");
ofstream fout("milkvisits.out");

struct qry {
  int req, lvl;
};

const int NMAX = 1e5 + 1;
int N, Q, a[NMAX], level[NMAX], ancestor[NMAX][18];
qry q[NMAX];
vector<int> G[NMAX], queries[NMAX], S[NMAX];
char sol[NMAX];

void dfs_lca(int u, int parent) {
  level[u] = level[parent] + 1;
  ancestor[u][0] = parent;
  for (int i = 1; i < 18; ++i)
    ancestor[u][i] = ancestor[ancestor[u][i - 1]][i - 1];
  for (int v : G[u])
    if (v != parent)
      dfs_lca(v, u);
}

int find_lca(int u, int v) {
  if (level[u] < level[v])
    swap(u, v);
  for (int lift = 17; lift >= 0; --lift)
    if (level[u] - (1 << lift) >= level[v])
      u = ancestor[u][lift];
  if (u == v)
    return u;
  for (int lift = 17; lift >= 0; --lift)
    if (ancestor[u][lift] != ancestor[v][lift]) {
      u = ancestor[u][lift];
      v = ancestor[v][lift];
    }
  return ancestor[u][0];
}

void dfs(int u, int parent) {
  S[a[u]].emplace_back(level[u]);
  for (int i : queries[u]) {
    int req = q[i].req;
    if (S[req].empty() || S[req].back() < q[i].lvl)
      continue;
    sol[i] = '1';
  }
  for (int v : G[u])
    if (v != parent)
      dfs(v, u);
  S[a[u]].pop_back();
}

void solve() {
  fin >> N >> Q;
  for (int i = 1; i <= N; ++i)
    fin >> a[i];
  for (int i = 1; i < N; ++i) {
    int u, v;
    fin >> u >> v;
    G[u].emplace_back(v);
    G[v].emplace_back(u);
  }
  dfs_lca(1, 0);
  for (int i = 0; i < Q; ++i) {
    int u, v;
    fin >> u >> v >> q[i].req;
    q[i].lvl = level[find_lca(u, v)];
    queries[u].emplace_back(i);
    queries[v].emplace_back(i);
    sol[i] = '0';
  }
  dfs(1, 0);
  for (int i = 0; i < Q; ++i)
    fout << sol[i];
}

void close_files() {
  fin.close();
  fout.close();
}

int main() {
  solve();
  close_files();
  return 0;
}