Problem p03392 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = (int)1e5 + 3;
const int infint = (int)1e9 + 3;
const int MOD = (int)998244353;
const ll inf = (ll)1e18 + 3;
string s;
unordered_map<string, int> M;
int dp[MAXN][3][3][3];
char nt(char a, char b) {
  if (a != 'a' && b != 'a')
    return 'a';
  if (a != 'b' && b != 'b')
    return 'b';
  return 'c';
}
void generate(string s) {
  M[s] = 1;
  for (int i = 0; i < s.size() - 1; i++)
    if (s[i] != s[i + 1]) {
      string t = s;
      t[i] = t[i + 1] = nt(s[i], s[i + 1]);
      if (M.count(t) == 0)
        generate(t);
    }
}
int add(int a, int b) {
  int c = a + b;
  if (c >= MOD)
    c -= MOD;
  if (c < 0)
    c += MOD;
  return c;
}
int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> s;

  if (s.size() <= 3) {
    generate(s);
    cout << M.size() << "\n";
    return 0;
  }
  bool diff = 0;
  for (int i = 1; i < s.size(); i++)
    if (s[i] != s[i - 1])
      diff = 1;
  if (!diff)
    return cout << 1, 0;
  int md = 0;
  for (int i = 0; i < s.size(); i++)
    if (s[i] == 'a')
      md += 0;
    else if (s[i] == 'b')
      md++;
    else
      md += 2;
  md %= 3;

  dp[1][0][0][0] = dp[1][1][1][0] = dp[1][2][2][0] = 1;
  for (int i = 2; i <= s.size(); i++)
    for (int md = 0; md < 3; md++)
      for (int last = 0; last < 3; last++) {
        for (int prev = 0; prev < 3; prev++) {
          if (prev != last)
            dp[i][md][last][0] = add(dp[i][md][last][0],
                                     dp[i - 1][(md - last + 3) % 3][prev][0]),
            dp[i][md][last][1] = add(dp[i - 1][(md - last + 3) % 3][prev][1],
                                     dp[i][md][last][1]);
          else
            dp[i][md][last][1] =
                add(add(dp[i][md][last][1],
                        dp[i - 1][(md - last + 3) % 3][prev][0]),
                    dp[i - 1][(md - last + 3) % 3][prev][1]);
        }
        // cerr << "# " << i << " " << md << " " << last << "-> " <<
        // dp[i][md][last][0] << " " << dp[i][md][last][1] << endl;
      }
  int ans = 0;
  for (int i = 0; i < 3; i++) {
    ans = add(ans, dp[s.size()][md][i][1]);
  }
  bool c = 0;
  for (int i = 1; i < s.size(); i++)
    if (s[i] == s[i - 1])
      c = 1;
  if (!c)
    ans = add(ans, 1);
  cout << ans;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = (int)2e5 + 3;
const int infint = (int)1e9 + 3;
const int MOD = (int)998244353;
const ll inf = (ll)1e18 + 3;
string s;
unordered_map<string, int> M;
int dp[MAXN][3][3][3];
char nt(char a, char b) {
  if (a != 'a' && b != 'a')
    return 'a';
  if (a != 'b' && b != 'b')
    return 'b';
  return 'c';
}
void generate(string s) {
  M[s] = 1;
  for (int i = 0; i < s.size() - 1; i++)
    if (s[i] != s[i + 1]) {
      string t = s;
      t[i] = t[i + 1] = nt(s[i], s[i + 1]);
      if (M.count(t) == 0)
        generate(t);
    }
}
int add(int a, int b) {
  int c = a + b;
  if (c >= MOD)
    c -= MOD;
  if (c < 0)
    c += MOD;
  return c;
}
int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);
  cin >> s;

  if (s.size() <= 3) {
    generate(s);
    cout << M.size() << "\n";
    return 0;
  }
  bool diff = 0;
  for (int i = 1; i < s.size(); i++)
    if (s[i] != s[i - 1])
      diff = 1;
  if (!diff)
    return cout << 1, 0;
  int md = 0;
  for (int i = 0; i < s.size(); i++)
    if (s[i] == 'a')
      md += 0;
    else if (s[i] == 'b')
      md++;
    else
      md += 2;
  md %= 3;

  dp[1][0][0][0] = dp[1][1][1][0] = dp[1][2][2][0] = 1;
  for (int i = 2; i <= s.size(); i++)
    for (int md = 0; md < 3; md++)
      for (int last = 0; last < 3; last++) {
        for (int prev = 0; prev < 3; prev++) {
          if (prev != last)
            dp[i][md][last][0] = add(dp[i][md][last][0],
                                     dp[i - 1][(md - last + 3) % 3][prev][0]),
            dp[i][md][last][1] = add(dp[i - 1][(md - last + 3) % 3][prev][1],
                                     dp[i][md][last][1]);
          else
            dp[i][md][last][1] =
                add(add(dp[i][md][last][1],
                        dp[i - 1][(md - last + 3) % 3][prev][0]),
                    dp[i - 1][(md - last + 3) % 3][prev][1]);
        }
        // cerr << "# " << i << " " << md << " " << last << "-> " <<
        // dp[i][md][last][0] << " " << dp[i][md][last][1] << endl;
      }
  int ans = 0;
  for (int i = 0; i < 3; i++) {
    ans = add(ans, dp[s.size()][md][i][1]);
  }
  bool c = 0;
  for (int i = 1; i < s.size(); i++)
    if (s[i] == s[i - 1])
      c = 1;
  if (!c)
    ans = add(ans, 1);
  cout << ans;
}