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

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

AtCoder ARC 094 F - Normalization (700 点)

これ系、この問題以降はたくさん見るようになったけど、この問題が先駆けとなった気がする

問題へのリンク

問題概要

'a', 'b', 'c' のみからなる長さ  N の文字列  S が与えられる。 S に以下の操作を好きな順序で好きな回数だけ行って得られる文字列が何通りあるか、998244353 で割ったあまりを求めよ。

制約

  •  2 \le N \le 2 \times 10^{5}

考えたこと

今となっては、こういう系の問題は mod. 3 かなにかで不変量を探すのがよい、という考察をするのが典型として広く知れ渡った。当時は多分そんなことはなかった。さて、a を 0、b を 1、c を 2 とすると、操作は

  • ab (01), ba (10) → cc (22)
  • bc (12), cb (21) -> aa (00)
  • ca (20), ac (02) -> bb (11)

のいずれかであって、すべて mod. 3 での合計値が不変であることがわかる。つまり、文字列  S の各文字の「総和 mod. 3」は不変量ということになる。この時点で、文字列  S と mod. 3 が異なる文字列は作れないことが確定する。

判定関数を考えて観察する

これもまた超典型手法だが、


操作によって出来上がるものが何通りあるかを問いかける問題では、まず、文字列 S, T が与えられて、S を T にすることができるかどうかの判定問題を解く


という風に考察すると上手くいくことが多い。さすがに、mod. 3 が等しかったらなんでもできるというわけにはいかない。極端な話、S = "aaaaa" とかに対しては何もできない。S = "aba" とかの場合は、できるのは

"aba",
"cca", "cbb", "aab", "acc", "bbc", "baa",
"acc", "bbc", "baa", "cca", "cbb", "aab"

の 13 個だけだ。mod. 3 が等しくても "cac" とかは作れないことに注意。つまり、"aab" -> "aba" のような変化はできない。こういう問題は、必要条件をひたすら列挙していって、いつか十分条件に到達するのを待つのがよいと思う。

必要条件を列挙していく

とりあえず、元の文字列 S を除くと、一度でも操作を行ったならば、「連続した 2 文字」が必ず含まれることになる。で、色々手を動かしていると、「S がすべて同じ文字であるケース」 と「S == T であるケース」はエッジケースとして、それ以外は

  • S と T の mod. 3 が等しい
  • T に連続した 2 文字が含まれる

という場合はほとんど可能なように思える。唯一の例外が S = "abc" のようなケース。こういうのは "aaa" と "ccc" しか作れない。

S の長さ N についての数学的帰納法で証明してみる。
N = 3 の場合は例外を含むので、N >= 4 について証明する。
N = 4 の場合はすべて可能なことは実験的に確かめる他なさそう (僕自身は下記の推論に基づいてサンプル通ったのできっとそういうことだと確信したのみでちゃんと調べてはいない...)。|S| = N (>4) の場合について、

  • T の連続する箇所が左端以外の場合は、S を上手に変形することで S[0] == T[0] とすることが可能なので、N -1 文字の場合に帰着 (T[1:N] も連続する箇所を含むので、帰納法の仮定から可能)
  • T の連続する箇所が左端の場合は、S を上手に変形することで S[N-1] == T[N-1] とすることが可能なので、N-1 文字の場合に帰着 (T[0:N-1] も連続する箇所を含むので、帰納法の仮定から可能)

というわけで、帰納法より

  • |S| >= 4
  • S = "aaaaaaaaaa" という形ではない
  • S と T の mod. 3 が等しい
  • T が連続する 2 文字を含む

という場合には常に可能であることが示された。

まとめ

結局、

  • 長さ N の 'a', 'b', 'c' の文字列であって
  • mod. 3 の値が所定の値であって
  • 「同じ文字が連続する箇所」を含む

ような文字列を数え上げらればよい。これはたとえば

  • dp[何文字目までみたか][総和 mod. 3][最後の文字][連続箇所があったか]

を更新する DP で求められる。これが求められた前提で、解答は以下の通り。


  • |S| = "aaaaa" の形のときは 1 通り
  • それ以外で |S| = 2 のときは 2 通り
  • それ以外で |S| = 3 のとき
    • S = "abc" のように 3 文字すべてが異なるときは 3 通り
    • S = "aba" のように両端が等しいときは 7 通り
    • S = "aab" のような場合は 6 通り
  • それ以外で |S| >= 4 のとき
    • S が「連続する箇所」を含むとき、DP で求めた値
    • S が「連続する箇所」を含まないとき、DP で求めた値 + 1

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

const int MOD = 998244353;
long long dp[210000][4][4][3]; // (mod 3, last, already)

void add(long long &a, long long b) {
    a += b;
    if (a >= MOD) a -= MOD;
}

long long solve(const string &S) {
    set<char> se;
    bool exist_ren = false;
    long long smod = 0;
    for (int i = 0; i < S.size(); ++i) {
        se.insert(S[i]);
        smod = (smod + (S[i] - 'a')) % 3;
        if (i + 1 < S.size() && S[i] == S[i+1]) exist_ren = true;
    }

    if (se.size() == 1) return 1;
    else if (S.size() == 2) return 2;
    else if (S.size() == 3) {
        if (se.size() == 3) return 3;
        else if (S[0] == S[2]) return 7;
        else return 6;
    }
    
    memset(dp, 0, sizeof(dp));
    for (int last = 0; last < 3; ++last) {
        dp[1][last][last][0] = 1;
    }
    for (int i = 1; i < S.size(); ++i) {
        for (int sum = 0; sum < 3; ++sum) {
            for (int last = 0; last < 3; ++last) {
                for (int last2 = 0; last2 < 3; ++last2) {
                    add(dp[i+1][(sum+last2)%3][last2][1], dp[i][sum][last][1]);
                    add(dp[i+1][(sum+last2)%3][last2][last==last2], dp[i][sum][last][0]);
                }
            }
        }
    }
    long long res = (exist_ren ? 0 : 1);
    for (int last = 0; last < 3; ++last) add(res, dp[S.size()][smod][last][1]);
    return res;
}

int main() {
    string S; cin >> S;
    cout << solve(S) << endl;
}

実験

りんごさんの解説を見ると、普通に実験した方が良さそうだったかもしれない。確かに |S| = 3 の場合の細かい場合分けはリスクが高かった...

https://www.youtube.com/watch?v=HDRfgn_UXLE