USACO 2019 December Contest, Gold Problem 2. Milk Visits - O(N + M)

#include <bits/stdc++.h>
#define pi pair<int,int>
#define f first
#define s second

using namespace std;

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

struct qry {
  int u, v, c;
};

const int NMAX = 1e5 + 1;
int N, Q, a[NMAX], label;
pi range[NMAX];
qry q[NMAX];
vector<int> G[NMAX], queries[NMAX], path;
vector<pi> S[NMAX];
char sol[NMAX];

void dfs1(int u, int parent) {
  range[u].f = ++label;
  for (int v : G[u])
    if (v != parent)
      dfs1(v, u);
  range[u].s = ++label;
}

bool is_ancestor(int x, int y) {
  return range[x].f <= range[y].f && range[y].s <= range[x].s;
}

void dfs2(int u, int parent) {
  S[a[u]].emplace_back(u, path.size());
  path.emplace_back(u);
  for (int i : queries[u]) {
    int c = q[i].c;
    if (S[c].empty())
      continue;
    pi v = S[c].back();
    if (v.f == u)
      sol[i] = '1';
    else {
      int node = path[v.s + 1];
      if (!is_ancestor(node, q[i].u + q[i].v - u))
        sol[i] = '1';
    }
  }
  for (int v : G[u])
    if (v != parent)
      dfs2(v, u);
  S[a[u]].pop_back();
  path.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);
  }
  dfs1(1, 0);
  for (int i = 0; i < Q; ++i) {
    fin >> q[i].u >> q[i].v >> q[i].c;
    queries[q[i].u].emplace_back(i);
    queries[q[i].v].emplace_back(i);
    sol[i] = '0';
  }
  dfs2(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;
}