JAG夏合宿2018 day3 D - Prefix Suffix Search

問題ページ

解法

単語のリストwをソートしたものを使うと二分探索で接頭辞がpの単語の集合がわかる。各単語を反転したリストw_invをソートした物を使うと二分探索で接尾辞がsの単語の集合がわかる。wとw_invをそれぞれ座標軸にとった二次元平面上にプロットすると各クエリは矩形中の点の個数を数えるクエリになる。二次元平面上の矩形中の点を数えるクエリは区間[l,r]でx以上y以下の要素の個数を求めるクエリ(rangefreq)になる。(参考 p44) rangefreqはセグ木を用いてO(log^2N)で求めることができる。(蟻本p170)

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
#define int ll
using PII = pair<int, int>;
template <typename T> using V = vector<T>;
template <typename T> using VV = vector<V<T>>;
template <typename T> using VVV = vector<VV<T>>;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()
#define PB push_back

const ll INF = (1LL<<60);
const int MOD = 1000000007;

template <typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template <typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; }
template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); }
template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
  out<<'('<<a.first<<','<<a.second<<')';
  return out;
}
template<class T>
ostream &operator <<(ostream& out,const vector<T>& a){
  out<<'[';
  REP(i, a.size()) {out<<a[i];if(i!=a.size()-1)out<<',';}
  out<<']';
  return out;
}

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0};

struct segTreeRangeFreq {
  int n;
  VV<int> dat;
  // aとbのマージを行う 
  V<int> f(V<int> a, V<int> b) {
    V<int> ret;
    int idx1 = 0, idx2 = 0;
    while(idx1 < a.size() || idx2 < b.size()) {
      if(idx2 >= b.size() || (idx1 < a.size() && a[idx1] < b[idx2])) ret.push_back(a[idx1]), idx1++;
      else ret.push_back(b[idx2]), idx2++;
    }
    return ret;
  }
  // 初期化 O(NlogN)
  segTreeRangeFreq() {}
  segTreeRangeFreq(V<int> v) {
    n = 1; while(n < v.size()) n *= 2;
    dat.assign(2*n-1, {});
    REP(i, v.size()) dat[i+n-1] = {v[i]};
    for(int i=n-2; i>=0; --i) dat[i] = f(dat[i*2+1], dat[i*2+2]);
  }
  // [a, b) のx以下の個数を返す O(log^2N)
  int query(int a, int b, int x, int k, int l, int r) {
    if(b <= l || r <= a) return 0;
    if(a <= l && r <= b) return upper_bound(ALL(dat[k]), x) - dat[k].begin();
    int vl = query(a, b, x, k*2+1, l, (l+r)/2);
    int vr = query(a, b, x, k*2+2, (l+r)/2, r);
    return vl + vr;
  }
  int query(int a, int b, int x) { return query(a, b, x, 0, 0, n); }
};

signed main(void)
{
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n, q;
  cin >> n >> q;
  vector<string> w(n), w_inv(n), s(q), p(q);
  vector<pair<string,int>> v(n);
  REP(i, n) cin >> w[i];
  REP(i, q) cin >> p[i] >> s[i], reverse(ALL(s[i]));
  
  sort(ALL(w));
  REP(i, n) {
    w_inv[i] = w[i];
    reverse(ALL(w_inv[i]));
    v[i] = {w_inv[i], i};
  }
  sort(ALL(w_inv));
  sort(ALL(v));

  V<int> init(n);
  REP(i, n) init[i] = v[i].second;
  segTreeRangeFreq seg(init);

  char c = 'z' + 1;

  REP(i, q) {
    // p[i], s[i] に対してのクエリ
    int x1 = lower_bound(ALL(w), p[i]) - w.begin();
    int x2 = lower_bound(ALL(w), p[i]+c) - w.begin() - 1;
    int y1 = lower_bound(ALL(w_inv), s[i]) - w_inv.begin();
    int y2 = lower_bound(ALL(w_inv), s[i]+c) - w_inv.begin() - 1;
    if(x2 < x1 || y2 < y1) {
      cout << 0 << endl;
      continue;
    }
 
    // 区間[y1,y2]に存在するx1以上x2以下の要素の数
    int vl = seg.query(y1, y2+1, x1-1);
    int vr = seg.query(y1, y2+1, x2);
    cout << vr - vl << endl;
  }

  return 0;
}