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

Apr 15th, 2021 (edited)
518
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <bits/stdc++.h>
2.
3. using namespace std;
4.
5. ifstream fin("milkvisits.in");
6. ofstream fout("milkvisits.out");
7.
8. struct qry {
9.   int req, lvl;
10. };
11.
12. const int NMAX = 1e5 + 1;
13. int N, Q, a[NMAX], level[NMAX], ancestor[NMAX][18];
14. qry q[NMAX];
15. vector<int> G[NMAX], queries[NMAX], S[NMAX];
16. char sol[NMAX];
17.
18. void dfs_lca(int u, int parent) {
19.   level[u] = level[parent] + 1;
20.   ancestor[u][0] = parent;
21.   for (int i = 1; i < 18; ++i)
22.     ancestor[u][i] = ancestor[ancestor[u][i - 1]][i - 1];
23.   for (int v : G[u])
24.     if (v != parent)
25.       dfs_lca(v, u);
26. }
27.
28. int find_lca(int u, int v) {
29.   if (level[u] < level[v])
30.     swap(u, v);
31.   for (int lift = 17; lift >= 0; --lift)
32.     if (level[u] - (1 << lift) >= level[v])
33.       u = ancestor[u][lift];
34.   if (u == v)
35.     return u;
36.   for (int lift = 17; lift >= 0; --lift)
37.     if (ancestor[u][lift] != ancestor[v][lift]) {
38.       u = ancestor[u][lift];
39.       v = ancestor[v][lift];
40.     }
41.   return ancestor[u][0];
42. }
43.
44. void dfs(int u, int parent) {
45.   S[a[u]].emplace_back(level[u]);
46.   for (int i : queries[u]) {
47.     int req = q[i].req;
48.     if (S[req].empty() || S[req].back() < q[i].lvl)
49.       continue;
50.     sol[i] = '1';
51.   }
52.   for (int v : G[u])
53.     if (v != parent)
54.       dfs(v, u);
55.   S[a[u]].pop_back();
56. }
57.
58. void solve() {
59.   fin >> N >> Q;
60.   for (int i = 1; i <= N; ++i)
61.     fin >> a[i];
62.   for (int i = 1; i < N; ++i) {
63.     int u, v;
64.     fin >> u >> v;
65.     G[u].emplace_back(v);
66.     G[v].emplace_back(u);
67.   }
68.   dfs_lca(1, 0);
69.   for (int i = 0; i < Q; ++i) {
70.     int u, v;
71.     fin >> u >> v >> q[i].req;
72.     q[i].lvl = level[find_lca(u, v)];
73.     queries[u].emplace_back(i);
74.     queries[v].emplace_back(i);
75.     sol[i] = '0';
76.   }
77.   dfs(1, 0);
78.   for (int i = 0; i < Q; ++i)
79.     fout << sol[i];
80. }
81.
82. void close_files() {
83.   fin.close();
84.   fout.close();
85. }
86.
87. int main() {
88.   solve();
89.   close_files();
90.   return 0;
91. }
92.
RAW Paste Data