Codeforces Round #296 (Div.1) D. Fuzzy Search

問題概要

文字列 S と T が与えられる.それぞれの文字数を n, m とする.
また,数 k が与えられるとする.
このとき,以下の条件をみたすような位置 p の数を求めよ.

すべての i (0 <= i < m) について,p+i-k <= x <= p+i+k を満たすある x が存在して T[i] == S[x] が成り立つ.

・制約
1 <= n, m <= 200000
0 <= k <= 200000
文字列 S および T は,'A','G','C','T'の文字からなる.

解法

問題を大雑把に理解するなら,各文字について多少のズレを許容したときに,S の中に T とマッチする場所はいくつあるか?ということ.
機械学習でCNNをやった人ならピンとくるかもしれないが,実は畳込みで実現できる.つまりFFTをやればいい.

どういうことか簡単に説明する.
S = AGCAATTCAT, T = ACAT, k = 1 とする.
まず,'A' だけに注目する.'A' が存在する位置と,それが影響する範囲(これはkのこと)に,1 を立てたビット列を S から生成する.
S -> 1111110111
T -> 1010 (T に関しては,影響させない)
S と T を畳み込んだとき,その値が T で1が立っている本数と一致すれば,Aに関する条件が満たされている,というわけである.
たとえば,p = 0 とする(AGCA と ACAT を比べる)と,値は 1*1 + 1*0 + 1*1 + 1*0 = 2 となり,条件を満たす.
しかし,p = 6 (TCAT と ACAT を比べる)とすると,0*1 + 1*0 + 1*1 + 1*0 = 1 となり,条件を満たさない.
これを 'A','G','C','T'について全部やって,総和が |T| と一致すれば,その場所は条件を満たす場所である.

ソースコード

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

// convolution は略. 

vector<int> calc(char c, string const& s, string const& t, int k) {
    int const n = s.size();
    int const m = t.size();
    vector<complex<double>> v(n);
    vector<complex<double>> w(m);
    for(int i = 0; i < n; ++i) {
        if(s[i] == c) {
            for(int x = i - 1; x >= max(0, i - k); --x) {
                if(real(v[x]) == 1) {
                    break;
                }
                v[x].real(1);
            }
            for(int x = i; x <= min(n - 1, i + k); ++x) {
                if(real(v[x]) == 1 || s[x] == c && x != i) {
                    break;
                }
                v[x].real(1);
            }
        }
    }
    for(int i = 0; i < m; ++i) {
        w[i].real(t[i] == c);
    }
    // FFT での畳込みは,v と w を逆方向にかけていく
    // ので,解法で述べたような計算をするには reverse する
    reverse(v.begin(), v.end());
    auto conv = convolution(v, w);
    vector<int> res(conv.size());
    for(int i = 0; i < (int)conv.size(); ++i) {
        res[i] = (int)(real(conv[i]) + 0.1);
    }
    return res;
}

int main() {
    int n, m, k;
    cin >> n >> m >> k;
    string s, t;
    cin >> s >> t;
    auto A = calc('A', s, t, k);
    auto T = calc('T', s, t, k);
    auto G = calc('G', s, t, k);
    auto C = calc('C', s, t, k);
    int res = 0;
    for(int i = 0; i < (int)A.size(); ++i) {
        res += A[i] + T[i] + G[i] + C[i] == m;
    }
    cout << res << endl;
}

感想

FFTの問題に慣れていなかったのもあって全然わからず,解説を聞いて驚きました.
個人的にはなかなか面白い問題でした.