Hatena::Grouptopcoder

tsubosakaの日記

2011-05-24

UTPC 2011 L

01:09

UTPCの解説とデータが上がってた(http://www.utpc.jp/2011/)ので、勉強も兼ねてWavelet Tree(もどき)を使った解法を書いてみた。手元だと最大ケースに大して4秒かかるので、ちょっとaccpetもらえるかは微妙なコードになった。

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;
struct BitVec{
  int N;
  vector<int> rank0;
  vector<int> rank0sign;
  void update(int pos , int diff , vector<int> & vec){
    if(pos > 0){
      vec[pos] = vec[pos - 1] + diff;
    }else{
      vec[pos] = diff;
    }
  }

  BitVec(vector<int> & arr , vector<int> & sign) 
    :N(arr.size()),rank0(arr.size()) , rank0sign(arr.size()) 
  {
    for(int i = 0 ; i < N ; ++i){
      int ai = arr[i];
      int si = sign[i];
      if(ai){
        update(i , 0 , rank0);
        update(i , 0 , rank0sign);
      }else{
        update(i , 1 , rank0);
        update(i , si , rank0sign);
      }
    }
  }
  //a[0..p]の中のbの数を返す
  int rank(int p , int b){
    if(b){
      return p + 1 - rank0[p];
    }else{
      return rank0[p];
    }
  }

  int rank0Sign(int p){
    return rank0sign[p];
  }
};

struct Tree{
  int hbit;
  BitVec * bv;
  struct Tree * left;
  struct Tree * right;
};

struct WTree{
  Tree * root;
  void gen(Tree * t , vector<int> & arr , vector<int> & sign){
    int hb = t->hbit;
    vector<int> bv(arr.size());
    vector<int> left, left_sign, right, right_sign;
    for(int i = 0 ; i < arr.size() ; ++i){
      if(arr[i] & hb){
        bv[i] = 1;
        right.push_back(arr[i]);
        right_sign.push_back(sign[i]);
      }else{
        bv[i] = 0;
        left.push_back(arr[i]);
        left_sign.push_back(sign[i]);
      }
    }
    t->bv = new BitVec(bv , sign);
    if(left.size() > 0 && hb){
      Tree * ltree = new Tree();
      ltree->hbit = hb / 2;
      gen(ltree , left , left_sign);
      t->left = ltree;
    }else{
      t->left = NULL;
    }
    if(right.size() > 0 && hb){
      Tree * rtree = new Tree();
      rtree->hbit = hb / 2;
      gen(rtree, right, right_sign);
      t->right = rtree;
    }else{
      t->right = NULL;
    }

  }
  WTree(vector<int> & arr , vector<int> & sign)
  {
    int m = *max_element(arr.begin() , arr.end());
    int hb = 1 << (31 - __builtin_clz(m));
    root = new Tree();
    root->hbit = hb;
    gen(root , arr , sign);
  }
  //a[0..p]のk以下のものの個数を返す
  int rank(int p , int k){
    Tree * r = root;
    int ret = 0;
    while(p >= 0 && r){
      if(k & r->hbit){
        ret += r->bv->rank0Sign(p);
        p = r->bv->rank(p , 1) - 1;
        r = r->right;
      }else{
        if(r->hbit == 0){
          ret += r->bv->rank0Sign(p);
        }
        p = r->bv->rank(p , 0) - 1;
        r = r->left;
      }
    }
    return ret;
  }
};

int pos;
void dfs(
    int cur , int parent, int level,
    vector<int> & seq, vector<int> & signs,
    vector<int> & levels,
    vector<int> & parents,
    vector<int> & firstHits,
    const vector<int> & values,
    const vector<pair<int , int> > & edges , const vector<int> & vstart)
{
  seq[pos] = values[cur]; 
  signs[pos] = 1;
  firstHits[cur] = pos;
  levels[cur] = level;
  parents[cur] = parent;
  ++pos;
  for(int i = vstart[cur] ; i < edges.size() && edges[i].first == cur ; ++i){
    int v = edges[i].second;
    if(v == parent)continue;
    dfs(v , cur , level + 1 , seq ,signs , levels , parents , firstHits , values , edges , vstart);
  }
  seq[pos] = values[cur]; 
  signs[pos] = -1;
  ++pos;
}

void build_lca(vector< vector<int> > & P , const vector<int> & parents){
  for(int i = 0 ; i < P.size() ; ++i)
    for(int j = 0 ; (1 << j) < P.size() ; ++j)
      P[i][j] = -1;

  for(int i = 0 ; i < P.size() ; ++i)
    P[i][0] = parents[i];

  for (int j = 1; (1 << j) < P.size(); j++)
    for (int i = 0; i < P.size(); i++)
      if (P[i][j - 1] != -1)
        P[i][j] = P[P[i][j - 1]][j - 1];
}

int lca(int p , int q, const vector< vector<int> > & P , const vector<int> & parents , const vector<int> & L){
  if(L[p] < L[q]){
    swap(p , q);
  }
  int log = 1;
  for (; (1 << log) <= L[p]; log++);
  log--;
  for (int i = log; i >= 0; i--)
    if (L[p] - (1 << i) >= L[q])
      p = P[p][i];
  if (p == q)
    return p;
  for (int i = log; i >= 0; i--)
    if (P[p][i] != -1 && P[p][i] != P[q][i])
      p = P[p][i], q = P[q][i];

  return parents[p];
}
int main(){
  int N , Q;
  cin >> N >> Q;
  vector<int> values(N);
  vector< pair<int , int> > edges(2 * (N - 1) );
  for(int i = 0 ; i < N ; ++i){
    cin >> values[i];
  }
  for(int i = 0 ; i < N - 1 ; ++i){
    int u , v;
    cin >> u >> v;
    --u; --v;
    edges[i * 2] = make_pair(u , v);
    edges[i * 2 + 1] = make_pair(v , u);
  }
  sort(edges.begin() , edges.end());
  vector<int> vstart(N , -1);
  for(int i = 0 ; i < edges.size() ; ++i){
    int u = edges[i].first;
    if(vstart[u] < 0){
      vstart[u] = i;
    }
  }
  pos = 0;
  vector<int> seq(2 * N);
  vector<int> signs(2 * N);
  vector<int> levels(N);
  vector<int> parents(N);
  vector<int> firstHits(N);
  dfs(0 , 0 , 0 , seq , signs , levels , parents, firstHits, values, edges , vstart);
  vector< vector<int> > P(N , vector<int>(32));
  build_lca(P , parents);
  WTree wt(seq , signs);
  sort(values.begin() , values.end());
  for(int i = 0 ; i < Q ; ++i){
    int v , w , l;
    cin >> v >> w >> l;
    --v; --w;
    int lvw = lca(v, w , P , parents , levels);
    int low = 0;
    int high = values.size() - 1; 
    while(low < high){
      int mid = low + (high - low) / 2;
      int mv  = values[mid];
      int num = wt.rank(firstHits[v] , mv) + wt.rank(firstHits[w] , mv) - wt.rank(firstHits[lvw] , mv);
      if(lvw){
        num -= wt.rank(firstHits[parents[lvw]] , mv);
      }
      if(num >= l){
        high = mid;
      }else{
        low = mid + 1;
      }
    }
    cout << values[low] << endl;
  }
  return 0;
}