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

Apr 15th, 2021
486
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <bits/stdc++.h>
2. #define pi pair<int,int>
3. #define f first
4. #define s second
5.
6. using namespace std;
7.
8. ifstream fin("milkvisits.in");
9. ofstream fout("milkvisits.out");
10.
11. struct qry {
12.   int u, v, c;
13. };
14.
15. const int NMAX = 1e5 + 1;
16. int N, Q, a[NMAX], label;
17. pi range[NMAX];
18. qry q[NMAX];
19. vector<int> G[NMAX], queries[NMAX], path;
20. vector<pi> S[NMAX];
21. char sol[NMAX];
22.
23. void dfs1(int u, int parent) {
24.   range[u].f = ++label;
25.   for (int v : G[u])
26.     if (v != parent)
27.       dfs1(v, u);
28.   range[u].s = ++label;
29. }
30.
31. bool is_ancestor(int x, int y) {
32.   return range[x].f <= range[y].f && range[y].s <= range[x].s;
33. }
34.
35. void dfs2(int u, int parent) {
36.   S[a[u]].emplace_back(u, path.size());
37.   path.emplace_back(u);
38.   for (int i : queries[u]) {
39.     int c = q[i].c;
40.     if (S[c].empty())
41.       continue;
42.     pi v = S[c].back();
43.     if (v.f == u)
44.       sol[i] = '1';
45.     else {
46.       int node = path[v.s + 1];
47.       if (!is_ancestor(node, q[i].u + q[i].v - u))
48.         sol[i] = '1';
49.     }
50.   }
51.   for (int v : G[u])
52.     if (v != parent)
53.       dfs2(v, u);
54.   S[a[u]].pop_back();
55.   path.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.   dfs1(1, 0);
69.   for (int i = 0; i < Q; ++i) {
70.     fin >> q[i].u >> q[i].v >> q[i].c;
71.     queries[q[i].u].emplace_back(i);
72.     queries[q[i].v].emplace_back(i);
73.     sol[i] = '0';
74.   }
75.   dfs2(1, 0);
76.   for (int i = 0; i < Q; ++i)
77.     fout << sol[i];
78. }
79.
80. void close_files() {
81.   fin.close();
82.   fout.close();
83. }
84.
85. int main() {
86.   solve();
87.   close_files();
88.   return 0;
89. }
90.
RAW Paste Data