Zero-Sum Ranges だった！！！

問題概要

'A', 'G', 'C', 'T' からなる文字列 $T$ が相補的であるとは、 $T$ を並び替えてできる文字列 $T'$ が存在して、

T[ i ] = 'A' ならば T'[ i ] = 'T'
T[ i ] = 'T' ならば T'[ i ] = 'A'
T[ i ] = 'G' ならば T'[ i ] = 'C'
T[ i ] = 'C' ならば T'[ i ] = 'G'

が成立することをいいます。

長さ $N$ の文字列 $S$ が与えられたとき、 $S$ の連続する部分文字列であって相補的なものが何通りあるか答えよ。

制約

$1 \le N \le 5000$

考えたこと

「相補的」という概念がこのままではわかりづらいので、わかりやすく言い換えよう。そうすると、次のことと同値であることがわかる。

文字列 $T$ に含まれる 'A' の個数と 'T' の個数とが等しく、'G' の個数と 'C' の個数とが等しい

このように言い換えれば、一気に扱いやすくなる。'A', 'G', 'C', 'T' をそれぞれ以下のように二次元ベクトルに置き換えよう

'A' = (1, 0)
'T' = (-1, 0)
'G' = (0, 1)
'C' = (0, -1)

このとき問題は「部分文字列中の二次元ベクトルの総和が (0, 0) になるものを数え上げよ」という問題と同じことになる。

累積和

ここまで来れば、有名な「Zero-Sum Ranges」とほとんど同じように解ける！

drken1215.hatenablog.com

実装を工夫すれば $O(N)$ になるだろうけど、下の実装だと $O(N \log N)$ になっている。ただ問題の制約的には $O(N^{2})$ が許されてるので、もう少し緩い方法でもいいかもしれない。

#include <bits/stdc++.h>

int main() {
    int N;
    string S;
    cin >> N >> S;
    using pint = pair<int,int>;
    map<pint, long long> ma;
    long long a = 0, b = 0;
    ma[pint(a, b)]++;
    for (int i = 0; i < N; ++i) {
        if (S[i] == 'A') ++a;
        else if (S[i] == 'T') --a;
        else if (S[i] == 'G') ++b;
        else --b;
        ma[pint(a,b)]++;
    }
    long long res = 0;
    for (auto it : ma) {
        long long n = it.second;
        res += n * (n-1) / 2;
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 104 B - DNA Sequence (茶色, 400 点)

問題概要

制約

考えたこと

累積和