AOJ 2614 Almost Same Substring

問題文

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

解法

ローリングハッシュでやります．

同じ長さの文字列 S, T に対して，「左からみて何文字一致しているか」と「右からみて何文字一致しているか」をそれぞれ二分探索で求めます．
もし S と T が一文字だけ異なるなら，左からの一致文字数と右からの一致文字数の和が，長さから1だけ小さいものになっているはずです．

よって，問題で与えられた S, T' について，S[a..a+|T'|-1] と T' に対し，先程やったことをすれば求まります．

ソースコード

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

// 1-indexed
// s[1...n]
class rolling_hash {
public:
    rolling_hash(std::string const& s)
        : n(s.size()),
          mod({999999937LL, 1000000007LL})
    {
        for(int i = 0; i < 2; ++i) {
            hs[i].resize(n + 1);
            p[i].resize(n + 1);
            hs[i][0] = 0;
            p[i][0] = 1;
            for(int j = 0; j < n; ++j) {
                hs[i][j + 1] = (hs[i][j] * base + s[j]) % mod[i];
                p[i][j + 1] = p[i][j] * base % mod[i];
            }
        }
    }

    // s[i...]
    long long query(int idx, int i) const {
        return hs[i][idx];
    }
    // s[l + 1...r]
    long long query(int l, int r, int i) const {
        return ((hs[i][r] - hs[i][l] * p[i][r - l]) % mod[i] + mod[i]) % mod[i];
    }

private:
    int n;
    std::vector<long long> hs[2];
    std::vector<long long> p[2];

    const std::array<long long, 2> mod;
    static const long long base = 29LL;
};

int main() {
    string S, T;
    cin >> S >> T;
    rolling_hash hs1(S), hs2(T);
    int res = 0;
    for(int i = 1; i + T.size() - 1 <= S.size(); ++i) {
        int lb = i - 1, ub = i + T.size();
        while(ub - lb > 1) {
            int m = (ub + lb) / 2;
            if(hs1.query(i - 1, m, 1) == hs2.query(0, m - i + 1, 1)) {
                lb = m;
            } else {
                ub = m;
            }
        }
        int l = lb;
        lb = i - 1, ub = i + T.size();
        while(ub - lb > 1) {
            int m = (ub + lb) / 2;
            if(hs1.query(m - 1, i + T.size() - 1, 1) == hs2.query(m - i, T.size(), 1)) {
                ub = m;
            } else {
                lb = m;
            }
        }
        int r = ub;

        res += l + 2 == r;
    }
    cout << res << endl;
}

感想

ローリングハッシュ使う時いっつもインデックスで混乱する．