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