Problem p02893 (C++)

Buggy submission (status: Runtime Error):
/* cerberus97 - Hanit Banga */

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>

using namespace std;

#define pb push_back
#define fast_cin()                                                             \
  ios_base::sync_with_stdio(false);                                            \
  cin.tie(NULL)

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int N = 2e5 + 10, mod = 998244353;

vector<pii> g[N];
int color[N];
ll cnt[N], p2[N];

ll solve(int d, int n, string x);
bool assign(int u, int col);

int main() {
  // int n = 6;
  // for (int i = 0; i < (1 << n); ++i) {
  // 	int cpy = i, steps = 0;
  // 	while (true) {
  // 		if (cpy % 2 == 0) {
  // 			cpy = (cpy / 2) + (1 << (n - 1));
  // 		} else {
  // 			cpy = (cpy - 1) / 2;
  // 		}
  // 		++steps;
  // 		if (cpy == i) {
  // 			break;
  // 		}
  // 	}
  // 	cout << i << ' ' << steps << endl;
  // }
  // return 0;
  fast_cin();
  int n;
  cin >> n;
  string x;
  cin >> x;
  // reverse(x.begin(), x.end());
  vector<int> divisors;
  int sq = sqrt(2 * n);
  for (int i = 1; i <= sq; ++i) {
    if ((2 * n) % i == 0) {
      divisors.pb(i);
      if (i * i != 2 * n) {
        divisors.pb((2 * n) / i);
      }
    }
  }
  p2[0] = 1;
  for (int i = 1; i < N; ++i) {
    p2[i] = (2 * p2[i - 1]) % mod;
  }
  sort(divisors.begin(), divisors.end());
  ll ans = 0;
  for (auto &d : divisors) {
    cnt[d] += solve(d, n, x);
    cnt[d] %= mod;
    // cout << d << ' ' << cnt[d] << endl;
    for (int i = 2 * d; i <= 2 * n; i += d) {
      cnt[i] += (mod - cnt[d]);
      cnt[i] %= mod;
    }
    ans += d * cnt[d];
    ans %= mod;
  }
  cout << ans << endl;
}

ll solve(int d, int n, string x) {
  for (int i = 0; i < n; ++i) {
    g[i].clear();
  }
  for (int i = 0; i < n; ++i) {
    if (i + d < n) {
      g[i].pb({i + d, 0});
      g[i + d].pb({i, 0});
    } else if (d != 2 * n) {
      g[i].pb({(i + d) % n, 1});
      g[(i + d) % n].pb({i, 1});
    } else {
      // g[i].pb({(i + d) % n, 0});
      // g[(i + d) % n].pb({i, 0});
    }
  }
  memset(color, -1, sizeof(color));
  ll components = 0;
  for (int i = 0; i < n; ++i) {
    if (color[i] == -1) {
      if (!assign(i, 0)) {
        return 0;
      }
      ++components;
    }
  }
  memset(color, -1, sizeof(color));
  // cout << "comp " << d << ' ' << components << endl;
  ll ans = 0, left = components;
  for (int i = 0; i < n; ++i) {
    // if (d == 2 * n) {
    // 	cout << i << ' ' << x[i] << ' ' << color[i] << ' ' << left << endl;
    // }
    if (x[i] == '0') {
      if (color[i] == 1) {
        break;
      } else if (color[i] == -1) {
        assign(i, 0);
        --left;
      }
    } else {
      if (color[i] == -1) {
        assign(i, 1);
        --left;
        ans += p2[left];
        ans %= mod;
      } else if (color[i] == 0) {
        ans += p2[left];
        ans %= mod;
        break;
      }
    }
    if (i == n - 1) {
      ++ans;
      ans %= mod;
    }
  }
  return ans;
}

bool assign(int u, int col) {
  if (color[u] != -1) {
    return color[u] == col;
  }
  color[u] = col;
  for (auto &e : g[u]) {
    int v = e.first, type = e.second;
    if (!assign(v, col ^ type)) {
      return false;
    }
  }
  return true;
}

Fixed submission:
/* cerberus97 - Hanit Banga */

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>

using namespace std;

#define pb push_back
#define fast_cin()                                                             \
  ios_base::sync_with_stdio(false);                                            \
  cin.tie(NULL)

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int N = 4e5 + 10, mod = 998244353;

vector<pii> g[N];
int color[N];
ll cnt[N], p2[N];

ll solve(int d, int n, string x);
bool assign(int u, int col);

int main() {
  // int n = 6;
  // for (int i = 0; i < (1 << n); ++i) {
  // 	int cpy = i, steps = 0;
  // 	while (true) {
  // 		if (cpy % 2 == 0) {
  // 			cpy = (cpy / 2) + (1 << (n - 1));
  // 		} else {
  // 			cpy = (cpy - 1) / 2;
  // 		}
  // 		++steps;
  // 		if (cpy == i) {
  // 			break;
  // 		}
  // 	}
  // 	cout << i << ' ' << steps << endl;
  // }
  // return 0;
  fast_cin();
  int n;
  cin >> n;
  string x;
  cin >> x;
  // reverse(x.begin(), x.end());
  vector<int> divisors;
  int sq = sqrt(2 * n);
  for (int i = 1; i <= sq; ++i) {
    if ((2 * n) % i == 0) {
      divisors.pb(i);
      if (i * i != 2 * n) {
        divisors.pb((2 * n) / i);
      }
    }
  }
  p2[0] = 1;
  for (int i = 1; i < N; ++i) {
    p2[i] = (2 * p2[i - 1]) % mod;
  }
  sort(divisors.begin(), divisors.end());
  ll ans = 0;
  for (auto &d : divisors) {
    cnt[d] += solve(d, n, x);
    cnt[d] %= mod;
    // cout << d << ' ' << cnt[d] << endl;
    for (int i = 2 * d; i <= 2 * n; i += d) {
      cnt[i] += (mod - cnt[d]);
      cnt[i] %= mod;
    }
    ans += d * cnt[d];
    ans %= mod;
  }
  cout << ans << endl;
}

ll solve(int d, int n, string x) {
  for (int i = 0; i < n; ++i) {
    g[i].clear();
  }
  for (int i = 0; i < n; ++i) {
    if (i + d < n) {
      g[i].pb({i + d, 0});
      g[i + d].pb({i, 0});
    } else if (d != 2 * n) {
      g[i].pb({(i + d) % n, 1});
      g[(i + d) % n].pb({i, 1});
    } else {
      // g[i].pb({(i + d) % n, 0});
      // g[(i + d) % n].pb({i, 0});
    }
  }
  memset(color, -1, sizeof(color));
  ll components = 0;
  for (int i = 0; i < n; ++i) {
    if (color[i] == -1) {
      if (!assign(i, 0)) {
        return 0;
      }
      ++components;
    }
  }
  memset(color, -1, sizeof(color));
  // cout << "comp " << d << ' ' << components << endl;
  ll ans = 0, left = components;
  for (int i = 0; i < n; ++i) {
    // if (d == 2 * n) {
    // 	cout << i << ' ' << x[i] << ' ' << color[i] << ' ' << left << endl;
    // }
    if (x[i] == '0') {
      if (color[i] == 1) {
        break;
      } else if (color[i] == -1) {
        assign(i, 0);
        --left;
      }
    } else {
      if (color[i] == -1) {
        assign(i, 1);
        --left;
        ans += p2[left];
        ans %= mod;
      } else if (color[i] == 0) {
        ans += p2[left];
        ans %= mod;
        break;
      }
    }
    if (i == n - 1) {
      ++ans;
      ans %= mod;
    }
  }
  return ans;
}

bool assign(int u, int col) {
  if (color[u] != -1) {
    return color[u] == col;
  }
  color[u] = col;
  for (auto &e : g[u]) {
    int v = e.first, type = e.second;
    if (!assign(v, col ^ type)) {
      return false;
    }
  }
  return true;
}