Advertisement
ivnikkk

Untitled

Feb 10th, 2023
956
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.09 KB | None | 0 0
  1. #define _CRT_SECURE_NO_WARNINGS
  2. #include "bits/stdc++.h"
  3. using namespace std;
  4. #define all(a) a.begin(), a.end()
  5. const int c = 400;
  6. vector<vector<pair<int, int>>> gr;
  7. vector<int> euler;
  8. struct Query {
  9.     int l, r, ind;
  10.     Query(int _l, int _r, int _i) : l(_l), r(_r), ind(_i) {}
  11.     Query() {}
  12.     bool operator<(const Query& other) {
  13.         if (l / c == other.l / c) {
  14.             if ((l / c) & 1) {
  15.                 return r > other.r;
  16.             }
  17.             return r < other.r;
  18.         }
  19.         return l < other.l;
  20.     }
  21. };
  22.  
  23. vector<Query> q;
  24. struct Set_sqrt {
  25.     int n;
  26.     vector<int> cnt, st, sq;
  27.     Set_sqrt(int _n) : n(_n) {
  28.         cnt.resize(_n + 1);
  29.         st.resize(_n + 1);
  30.         sq.resize(c + 1);
  31.     }
  32.     void insert(int x) {
  33.         cnt[x]++;
  34.         if (cnt[x] == 2) {
  35.             sq[x / c]++;
  36.             st[x] = 1;
  37.         }
  38.     }
  39.     void erase(int x) {
  40.         cnt[x]--;
  41.         if (cnt[x] == 1) {
  42.             sq[x / c]--;
  43.             st[x] = 0;
  44.         }
  45.     }
  46.     int get_max() {
  47.         int it = c;
  48.         while (sq[it] == 0) {
  49.             it--;
  50.             if (it == -1) {
  51.                 return 0;
  52.             }
  53.         }
  54.         int res = min(it * c + c - 1, n);
  55.         while (st[res] == 0) {
  56.             res--;
  57.         }
  58.         return res;
  59.     }
  60. };
  61. vector<int> clr;
  62. void dfs(int v, int p) {
  63.     euler.push_back(clr[v]);
  64.     for (auto&& [u, ind] : gr[v]) {
  65.         if (u == p) { continue; }
  66.         int l = (int)euler.size();
  67.         dfs(u, v);
  68.         int r = (int)euler.size() - 1;
  69.         q.push_back(Query(l, r, ind));
  70.     }
  71. }
  72. signed main() {
  73. #ifdef _DEBUG
  74.     freopen("input.txt", "r", stdin);
  75.     freopen("output.txt", "w", stdout);
  76. #endif
  77.     ios_base::sync_with_stdio(false);
  78.     cin.tie(nullptr);
  79.     int n; cin >> n;
  80.     gr.resize(n); clr.resize(n);
  81.     for (int i = 0; i < n - 1; i++) {
  82.         int u, v; cin >> u >> v;
  83.         u--, v--;
  84.         gr[u].push_back({ v, i });
  85.         gr[v].push_back({ u, i });
  86.     }
  87.     for (int i = 0; i < n; i++) {
  88.         cin >> clr[i];
  89.     }
  90.     vector<int> cnt1;
  91.     vector<int> cnt2;
  92.     vector<int> cpy = clr;
  93.     cpy.push_back(0);
  94.     sort(all(cpy));
  95.     cpy.resize(unique(all(cpy)) - cpy.begin());
  96.     Set_sqrt U1((int)cpy.size());
  97.     Set_sqrt U2((int)cpy.size());
  98.  
  99.     for (int& i : clr) {
  100.         i = lower_bound(all(cpy), i) - cpy.begin();
  101.         U1.insert(i);
  102.     }
  103.     dfs(0, 0);
  104.     sort(all(q));
  105.     auto add = [&](int x) {
  106.         U1.erase(x);
  107.         U2.insert(x);
  108.     };
  109.     auto del = [&](int x) {
  110.         U1.insert(x);
  111.         U2.erase(x);
  112.     };
  113.     vector<int> ans(n - 1);
  114.     int l = 0, r = -1;
  115.     for (const Query& i : q) {
  116.         while (l > i.l) {
  117.             add(euler[--l]);
  118.         }
  119.         while (r < i.r) {
  120.             add(euler[++r]);
  121.         }
  122.         while (l < i.l) {
  123.             del(euler[l++]);
  124.         }
  125.         while (r > i.r) {
  126.             del(euler[r--]);
  127.         }
  128.         ans[i.ind] = cpy[max(U1.get_max(), U2.get_max())];
  129.     }
  130.     for (int& i : ans) {
  131.         cout << i << '\n';
  132.     }
  133.     return 0;
  134. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement