Guest User

Untitled

a guest
Feb 2nd, 2017
210
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.66 KB | None | 0 0
  1. ArrayList<Integer> graph[];
  2. int[] c;
  3. int[] parentColor;
  4.  
  5. void solve() throws IOException {
  6. int n = rI();
  7. this.graph = new ArrayList[n];
  8. for (int i = 0; i < n; i++) {
  9. graph[i] = new ArrayList<Integer>();
  10. }
  11. for (int i = 0; i < n - 1; i++) {
  12. int a = rI() - 1;
  13. int b = rI() - 1;
  14. graph[a].add(b);
  15. graph[b].add(a);
  16. }
  17. this.parentColor = new int[n];
  18. this.c = rA(n);
  19. this.used = new boolean[n];
  20. this.listForCheckSecond = new ArrayList<Integer>();
  21. this.listForCheckFirst = new ArrayList<Integer>();
  22. this.edge = 0;
  23. this.firstColor = c[0];
  24. this.secondColor = -2;
  25. this.f = false;
  26. dfs(0);
  27. Arrays.fill(used, false);
  28. dfs2(edge);
  29. if (f) {
  30. Arrays.fill(used, false);
  31. edge = 0;
  32. f = false;
  33. dfs2(0);
  34. if (f) {
  35. out.print("NO");
  36. } else {
  37. out.println("YES");
  38. out.println("1");
  39. }
  40. } else {
  41. out.println("YES");
  42. out.println(edge + 1);
  43. }
  44.  
  45. }
  46.  
  47. int edge;
  48. boolean[] used;
  49. ArrayList<Integer> listForCheckSecond;
  50. ArrayList<Integer> listForCheckFirst;
  51. int firstColor;
  52. int secondColor;
  53. boolean f;
  54.  
  55. // =======================================================
  56. void dfs(int from) {
  57. used[from] = true;
  58. for (int i : graph[from]) {
  59. if (!used[i]) {
  60. if (c[i] != firstColor) {
  61. if (edge == 0) {
  62. edge = i;
  63. continue;
  64. } else {
  65. f = true;
  66. return;
  67. }
  68. } else {
  69. dfs(i);
  70. }
  71.  
  72. }
  73. }
  74. }
  75.  
  76. void dfs2(int from) {
  77. used[from] = true;
  78. for (int i : graph[from]) {
  79. if (!used[i]) {
  80. if (from == edge || c[from] == c[i]) {
  81. dfs2(i);
  82. } else {
  83. f = true;
  84. return;
  85. }
  86. }
  87. }
  88. }
Advertisement
Add Comment
Please, Sign In to add comment