StormFalcon32

MilkVisits

Dec 29th, 2019
139
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 6.33 KB | None | 0 0
  1. import java.io.BufferedReader;
  2. import java.io.BufferedWriter;
  3. import java.io.FileReader;
  4. import java.io.FileWriter;
  5. import java.io.IOException;
  6. import java.io.PrintWriter;
  7. import java.util.ArrayList;
  8. import java.util.LinkedList;
  9. import java.util.Stack;
  10. import java.util.StringTokenizer;
  11.  
  12. public class MilkVisits {
  13.  
  14. static int[] depth;
  15. static int[][] parent;
  16. static int level;
  17. static LinkedList<Integer>[] adjList;
  18. static int[] cows;
  19. static int[][] query;
  20. static LinkedList<Integer>[] queryList;
  21. static ArrayList<Integer>[] lastCows;
  22. static boolean[] ans;
  23. static int[] preorderNums;
  24. static int[] postorderNums;
  25. static int N;
  26.  
  27. @SuppressWarnings("unchecked")
  28. public static void main(String[] args) throws IOException {
  29. BufferedReader in = new BufferedReader(new FileReader("D:\\Java\\USACO-Gold\\Gold\\MilkVisits\\4.in"));
  30. // BufferedReader in = new BufferedReader(new FileReader("milkvisits.in"));
  31. PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter("milkvisits.out")));
  32. StringTokenizer tk = new StringTokenizer(in.readLine());
  33. // input vars
  34. N = Integer.parseInt(tk.nextToken());
  35. int M = Integer.parseInt(tk.nextToken());
  36. cows = new int[N];
  37. // lca vars
  38. level = (int) (Math.ceil(Math.log(N) / Math.log(2)));
  39. depth = new int[N];
  40. parent = new int[N][level];
  41. // other vars
  42. adjList = new LinkedList[N];
  43. queryList = new LinkedList[N];
  44. lastCows = new ArrayList[N];
  45. preorderNums = new int[N];
  46. postorderNums = new int[N];
  47. // ans vars
  48. ans = new boolean[M];
  49. for (int i = 0; i < N; i++) {
  50. adjList[i] = new LinkedList<Integer>();
  51. queryList[i] = new LinkedList<Integer>();
  52. lastCows[i] = new ArrayList<Integer>();
  53. }
  54. tk = new StringTokenizer(in.readLine());
  55. for (int i = 0; i < N; i++) {
  56. cows[i] = Integer.parseInt(tk.nextToken()) - 1;
  57. }
  58. for (int i = 0; i < N - 1; i++) {
  59. tk = new StringTokenizer(in.readLine());
  60. int x = Integer.parseInt(tk.nextToken()) - 1;
  61. int y = Integer.parseInt(tk.nextToken()) - 1;
  62. adjList[x].add(y);
  63. adjList[y].add(x);
  64. }
  65. dfsLCA(0, -1);
  66. precomputeSparseMatrix(N);
  67. query = new int[M][3];
  68. for (int i = 0; i < M; i++) {
  69. tk = new StringTokenizer(in.readLine());
  70. query[i][0] = Integer.parseInt(tk.nextToken()) - 1;
  71. query[i][1] = Integer.parseInt(tk.nextToken()) - 1;
  72. query[i][2] = Integer.parseInt(tk.nextToken()) - 1;
  73. queryList[query[i][0]].add(i);
  74. queryList[query[i][1]].add(i);
  75. }
  76. preandpost();
  77. dfs(0, -1, 0);
  78. for (int i = 0; i < M; i++) {
  79. out.print(ans[i] ? "1" : "0");
  80. }
  81. out.println();
  82. out.close();
  83. in.close();
  84. }
  85.  
  86. static void dfs(int curr, int prev, int depth) {
  87. lastCows[cows[curr]].add(curr);
  88. for (int i : queryList[curr]) {
  89. int a = query[i][0];
  90. int b = query[i][1];
  91. int c = query[i][2];
  92. if (!lastCows[c].isEmpty()) {
  93. int last = lastCows[c].get(lastCows[c].size() - 1);
  94. if (last == a || last == b) {
  95. ans[i] = true;
  96. }
  97. if (!ancestor(last, b) || !ancestor(last, a)) {
  98. ans[i] = true;
  99. }
  100. if (lca(a, b) == last) {
  101. ans[i] = true;
  102. }
  103. }
  104. }
  105. for (int adj : adjList[curr]) {
  106. if (adj != prev) {
  107. dfs(adj, curr, depth++);
  108. }
  109. }
  110. lastCows[cows[curr]].remove(lastCows[cows[curr]].size() - 1);
  111. }
  112.  
  113. static void preandpost() {
  114. int num = 0;
  115. boolean[] visited = new boolean[N];
  116. Stack<Integer> s1 = new Stack<Integer>();
  117. s1.push(0);
  118. while (!s1.isEmpty()) {
  119. int curr = s1.pop();
  120. visited[curr] = true;
  121. preorderNums[curr] = num;
  122. num++;
  123. for (int adj : adjList[curr]) {
  124. if (!visited[adj]) {
  125. s1.push(adj);
  126. }
  127. }
  128. }
  129. s1 = new Stack<Integer>();
  130. Stack<Integer> s2 = new Stack<Integer>();
  131. visited = new boolean[N];
  132. s1.add(0);
  133. while (!s1.isEmpty()) {
  134. int curr = s1.pop();
  135. visited[curr] = true;
  136. s2.push(curr);
  137. for (int adj : adjList[curr]) {
  138. if (!visited[adj]) {
  139. s1.push(adj);
  140. }
  141. }
  142. }
  143. num = 0;
  144. while (!s2.isEmpty()) {
  145. int curr = s2.pop();
  146. postorderNums[curr] = num;
  147. num++;
  148. }
  149. }
  150.  
  151. static boolean ancestor(int a, int b) {
  152. return preorderNums[a] <= preorderNums[b] && postorderNums[b] <= postorderNums[a];
  153. }
  154.  
  155. static void dfsLCA(int curr, int prev) {
  156. if (curr != 0) {
  157. depth[curr] = depth[prev] + 1;
  158. parent[curr][0] = prev;
  159. }
  160. for (int adj : adjList[curr]) {
  161. if (adj != prev) {
  162. dfsLCA(adj, curr);
  163. }
  164. }
  165. }
  166.  
  167. static void precomputeSparseMatrix(int N) {
  168. for (int i = 1; i < level; i++) {
  169. for (int node = 0; node < N; node++) {
  170. if (parent[node][i - 1] != -1)
  171. parent[node][i] = parent[parent[node][i - 1]][i - 1];
  172. }
  173. }
  174. }
  175.  
  176. static int lca(int u, int v) {
  177. if (depth[v] < depth[u]) {
  178. int temp = u;
  179. u = v;
  180. v = temp;
  181. }
  182. int diff = depth[v] - depth[u];
  183. for (int i = 0; i < level; i++) {
  184. if (((diff >> i) & 1) == 1) {
  185. v = parent[v][i];
  186. }
  187. }
  188. if (u == v) {
  189. return u;
  190. }
  191. for (int i = level - 1; i >= 0; i--) {
  192. if (parent[u][i] != parent[v][i]) {
  193. u = parent[u][i];
  194. v = parent[v][i];
  195. }
  196. }
  197. return parent[u][0];
  198. }
  199. }
Advertisement
Add Comment
Please, Sign In to add comment