AOJ 2907 Prefix Suffix Search

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2907

解法

\(w_i\) をそれぞれ逆にした文字列集合を \(\{rw_i\}\) とする．
最初にそれぞれをもともとのインデックスが分かるようにソートしておく．
以降は \(w_i, rw_i\) はソート済みであるとする．

各 \(i\) に対して，\(w_i\) のもともとのインデックスに対応する文字列が \(rw_j\) であるとする．
すると，\((0, j_0), (1, j_1), \ldots, (n, j_n)\) という2次元上の点列が得られる．

各クエリ \(p, s\) に対して，\(p\) （あるいは \(s\)）の条件を満たす文字列集合は，\(w_i\) （あるいは \(rw_i\)）における1つの連続区間とみなすことができる．
\(p, s\) に対応する区間をそれぞれ \([l_1, r_1], [l_2, r_2]\) とする．
すると，先の2次元上の点列と合わせて，以下の値を求められれば良い．

\[\# \{ j_i ~| ~l_1 \leq i \leq r_1, l_2 \leq j_i \leq r_2\}\]

これは WaveletMatrix やマージソートをセグメント木っぽく行うデータ構造などで高速に処理することができる．ICPC など写経が必要なら計算量は不利だが後者を選ぶとよいだろう．

また，\(p\) や \(s\) に対する区間を求めるのは単純に二分探索すれば計算量的にも問題ない．

ソースコード

#include <bits/stdc++.h>
using namespace std;


template <typename T>
class merge_tree {
public:
    merge_tree(std::vector<T> const& v)
        : n(size(v.size())), dat(n * 2)
    {
        for(int i = 0; i < (int)v.size(); ++i) {
            dat[i + n].push_back(v[i]);
        }
        for(int i = n - 1; i >= 0; --i) {
            std::merge(dat[i * 2].begin(), dat[i * 2].end(),
                       dat[i * 2 + 1].begin(), dat[i * 2 + 1].end(),
                       std::back_inserter(dat[i]));
        }
    }

    // in [l, r), the number of values larger than lb
    int count(int l, int r, T lb) {
        l += n, r += n;
        int res = 0;
        while(l < r) {
            if(l & 1) res += count(l++, lb);
            if(r & 1) res += count(--r, lb);
            l >>= 1, r >>= 1;
        }
        return res;
    }

private:
    int size(int x) {
        int res = 1;
        while(res < x) res <<= 1;
        return res;
    }

    int count(int node_id, T lb) {
        return std::lower_bound(dat[node_id].begin(), dat[node_id].end(), lb) - dat[node_id].begin();
    }

private:
    const int n;
    std::vector<std::vector<T>> dat;
};


int main() {
    using ps = pair<string, int>;

    int n, q; cin >> n >> q;
    vector<ps> ws(n), rev_ws(n);
    for(int i = 0; i < n; ++i) {
        cin >> ws[i].first;
        rev_ws[i] = ws[i];
        ws[i].second = rev_ws[i].second = i;
        reverse(rev_ws[i].first.begin(), rev_ws[i].first.end());
    }
    sort(ws.begin(), ws.end());
    sort(rev_ws.begin(), rev_ws.end());

    vector<int> ids(n);
    { // build ids
        vector<int> rev_ids(n);
        for(int i = 0; i < n; ++i) {
            rev_ids[rev_ws[i].second] = i;
        }
        for(int i = 0; i < n; ++i) {
            ids[i] = rev_ids[ws[i].second];
        }
    }
    merge_tree<int> mt(ids);

    auto calc_range = [] (vector<ps> const& bag, string const& s) -> pair<int, int> {
        int lb = -1, ub = bag.size();
        while(ub - lb > 1) {
            const int mid = (lb + ub) / 2;
            (bag[mid].first.compare(0, s.size(), s) < 0 ? lb : ub) = mid;
        }
        const int res_l = ub;
        lb = -1, ub = bag.size();
        while(ub - lb > 1) {
            const int mid = (lb + ub) / 2;
            (bag[mid].first.compare(0, s.size(), s) <= 0 ? lb : ub) = mid;
        }
        return {res_l, ub};
    };

    while(q--) {
        string prefix, suffix; cin >> prefix >> suffix;
        reverse(suffix.begin(), suffix.end());
        const auto [l1, r1] = calc_range(ws, prefix);
        const auto [l2, r2] = calc_range(rev_ws, suffix);
        cout << mt.count(l1, r1, r2) - mt.count(l1, r1, l2) << endl;
    }
}

感想

久々に AOJ を使ったら C++17 が使えるようになっていてよかった．