SPOJ/COT.cc

   1 #include <cstdio>
   2 #include <queue>
   3 #include <vector>
   4 #include <algorithm>
   5
   6 using namespace std;
   7
   8 const int maxn = 1e5 + 10;
   9
  10 int n, m, w[maxn], rlim;
  11 vector<int> table;
  12
  13 vector<int> epool[maxn];
  14
  15 struct node {
  16   node *ch[2];
  17   int size;
  18 } pool[maxn << 6], *nil = pool, *tot = pool, *root[maxn];
  19
  20 int dp[maxn][18], dep[maxn]; //lca
  21
  22 node* Modify(node *orig, int l, int r, int pos) {
  23   node *cur = ++tot;
  24   *cur = *orig;
  25   cur->size++;
  26   if (l == r) return cur;
  27   int mid = (l + r) >> 1;
  28   if (pos <= mid)
  29     cur->ch[0] = Modify(orig->ch[0], l, mid, pos);
  30   else
  31     cur->ch[1] = Modify(orig->ch[1], mid + 1, r, pos);
  32   return cur;
  33 }
  34
  35 void turn(vector<node*> &vec, int d) {
  36   vector<node*>::iterator iter;
  37   for (iter = vec.begin(); iter != vec.end(); ++iter)
  38     if (*iter) *iter = (*iter)->ch[d];
  39 }
  40
  41 int calc(vector<node*> &vec) {
  42   int res = 0;
  43   vector<node*>::iterator iter;
  44   for (iter = vec.begin(); iter != vec.end(); ++iter)
  45     if (*iter && (*iter)->ch[0]) res += (*iter)->ch[0]->size;
  46   return res;
  47 }
  48
  49 int Query(int l, int r, int k, vector<node*> &add, vector<node*> &del) {
  50   if (l == r) return l;
  51   int sw = calc(add) - calc(del);
  52   int mid = (l + r) >> 1;
  53   bool d = k > sw;
  54   turn(add, d);
  55   turn(del, d);
  56   if (d) {
  57     k -= sw;
  58     return Query(mid + 1, r, k, add, del);
  59   } else {
  60     return Query(l, mid, k, add, del);
  61   }
  62 }
  63
  64 void bfs(int s) {
  65   queue<int> q;
  66   q.push(s);
  67   root[s] = Modify(nil, 0, rlim, w[s]);
  68   dp[s][0] = s;
  69   dep[s] = 1;
  70   vector<int>::iterator iter;
  71   while (!q.empty()) {
  72     int a = q.front();
  73     q.pop();
  74     for (iter = epool[a].begin(); iter != epool[a].end(); ++iter) {
  75       if (!root[*iter]) {
  76         root[*iter] = Modify(root[a], 0, rlim, w[*iter]);
  77         q.push(*iter);
  78         dep[*iter] = dep[a] + 1;
  79         dp[*iter][0] = a;
  80       }
  81     }
  82   }
  83 }
  84
  85 int lca(int u, int v) {
  86   if (dep[u] < dep[v]) swap(u, v);
  87   for (int i = 17; i >= 0; --i) {
  88     if (dep[dp[u][i]] >= dep[v]) {
  89       u = dp[u][i];
  90       if (dep[u] == dep[v]) break;
  91     }
  92   }
  93   if (u == v) return u;
  94   for (int i = 17; i >= 0; --i)
  95     if (dp[u][i] != dp[v][i])
  96       u = dp[u][i], v = dp[v][i];
  97   return dp[u][0];
  98 }
  99
 100 int query(int u, int v, int k) {
 101   int p = lca(u, v);
 102   vector<node*> add, del;
 103   add.push_back(root[u]);
 104   add.push_back(root[v]);
 105   del.push_back(root[p]);
 106   if (dp[p][0] != p) del.push_back(root[dp[p][0]]);
 107   return table[Query(0, rlim, k, add, del)];
 108 }
 109
 110 int main() {
 111   int n, m;
 112   scanf("%d%d", &n, &m);
 113   for (int i = 1; i <= n; ++i) scanf("%d", w + i), table.push_back(w[i]);
 114   sort(table.begin(), table.end());
 115   vector<int>::iterator iter = unique(table.begin(), table.end());
 116   table.erase(iter, table.end());
 117   rlim = table.size();
 118   for (int i = 1; i <= n; ++i)
 119     w[i] = lower_bound(table.begin(), table.end(), w[i]) - table.begin();
 120   int u, v, k;
 121   for (int i = 1; i < n; ++i) {
 122     scanf("%d%d", &u, &v);
 123     epool[u].push_back(v);
 124     epool[v].push_back(u);
 125   }
 126   nil->ch[0] = nil->ch[1] = nil;
 127   bfs(1);
 128   for (int j = 1; j < 18; ++j)
 129     for (int i = 1; i <= n; ++i)
 130       dp[i][j] = dp[dp[i][j-1]][j-1];
 131   while (m--) {
 132     scanf("%d%d%d", &u, &v, &k);
 133     printf("%d\n", query(u, v, k));
 134   }
 135   delete nil;
 136   return 0;
 137 }