Problem p02609 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define FOR_LT(i, beg, end) for (int i = (int)(beg); i < (int)(end); i++)
#define FOR_LE(i, beg, end) for (int i = (int)(beg); i <= (int)(end); i++)
#define FOR_DW(i, beg, end) for (int i = (int)(beg); (int)(end) <= i; i--)
#define REP(n) for (int repeat_index = 0; repeat_index < (int)n; repeat_index++)

using namespace std;

int64_t pow_rem(int64_t x, int64_t m, int64_t rem) {
  if (m == 0) {
    return 1;
  }
  if (m == 1) {
    return x;
  }

  if (m & 1) {
    int64_t val = pow_rem(x, m / 2, rem);
    val = val * val;
    val %= rem;
    val *= x;
    val %= rem;
    return val;
  } else {
    int64_t val = pow_rem(x, m / 2, rem);
    val = val * val;
    val %= rem;
    return val;
  }
}

int64_t solve(int64_t val) {
  if (val == 0)
    return 0;

  static unordered_map<int64_t, int64_t> memo;
  auto it = memo.find(val);
  if (it != memo.end()) {
    return it->second;
  }

  int64_t cnt = 0;
  FOR_LT(i, 0, 32) {
    if (val & (1 << i))
      cnt++;
  }
  int64_t nv = val % cnt;
  int64_t ans = solve(nv) + 1;
  memo.insert({val, ans});
  return ans;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout << fixed << setprecision(20);

  int n;
  cin >> n;
  string s;
  cin >> s;

  int cnt = 0;
  FOR_LT(i, 0, s.size()) {
    if (s[i] == '1')
      cnt++;
  }

  int64_t nu = 0, nd = 0;
  int64_t cntu = cnt + 1;
  int64_t cntd = cnt - 1;
  int64_t vu = 1, vd = 1;
  vector<int64_t> vus(n + 1);
  vector<int64_t> vds(n + 1);
  FOR_DW(i, n - 1, 0) {
    vus[i] = vu;
    vds[i] = vd;
    if (s[i] == '1') {
      nu += vu;
      nd += vd;
    }
    vu *= 2;
    vu %= cntu;
    if (cntd >= 1) {
      vd *= 2;
      vd %= cntd;
    }
  }

  FOR_LT(i, 0, n) {
    if (s[i] == '1') {
      if (cntd <= 0) {
        cout << 0 << endl;
        continue;
      }
      int64_t cv = nd - vds[i];
      cv %= cntd;
      if (cv < 0)
        cv += cntd;
      // cout << i << " " << nd << " " << cv << " " << cntd << endl;
      cout << solve(cv) + 1 << endl;
    } else {
      int64_t cv = nu + vus[i];
      cv %= cntu;
      // cout << i << " " << nu << " " << cv << " " << cntu << endl;
      cout << solve(cv) + 1 << endl;
    }
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define FOR_LT(i, beg, end) for (int i = (int)(beg); i < (int)(end); i++)
#define FOR_LE(i, beg, end) for (int i = (int)(beg); i <= (int)(end); i++)
#define FOR_DW(i, beg, end) for (int i = (int)(beg); (int)(end) <= i; i--)
#define REP(n) for (int repeat_index = 0; repeat_index < (int)n; repeat_index++)

using namespace std;

int64_t pow_rem(int64_t x, int64_t m, int64_t rem) {
  if (m == 0) {
    return 1;
  }
  if (m == 1) {
    return x;
  }

  if (m & 1) {
    int64_t val = pow_rem(x, m / 2, rem);
    val = val * val;
    val %= rem;
    val *= x;
    val %= rem;
    return val;
  } else {
    int64_t val = pow_rem(x, m / 2, rem);
    val = val * val;
    val %= rem;
    return val;
  }
}

int64_t solve(int64_t val) {
  if (val == 0)
    return 0;

  static unordered_map<int64_t, int64_t> memo;
  auto it = memo.find(val);
  if (it != memo.end()) {
    return it->second;
  }

  int64_t cnt = 0;
  FOR_LT(i, 0, 32) {
    if (val & (1 << i))
      cnt++;
  }
  int64_t nv = val % cnt;
  int64_t ans = solve(nv) + 1;
  memo.insert({val, ans});
  return ans;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout << fixed << setprecision(20);

  int n;
  cin >> n;
  string s;
  cin >> s;

  int cnt = 0;
  FOR_LT(i, 0, s.size()) {
    if (s[i] == '1')
      cnt++;
  }

  int64_t nu = 0, nd = 0;
  int64_t cntu = cnt + 1;
  int64_t cntd = cnt - 1;
  int64_t vu = 1, vd = 1;
  vector<int64_t> vus(n + 1);
  vector<int64_t> vds(n + 1);
  FOR_DW(i, n - 1, 0) {
    vus[i] = vu;
    vds[i] = vd;
    if (s[i] == '1') {
      nu += vu;
      nd += vd;
    }
    vu *= 2;
    vu %= cntu;
    if (cntd >= 1) {
      vd *= 2;
      vd %= cntd;
    }
  }

  FOR_LT(i, 0, n) {
    if (s[i] == '1') {
      if (cntd <= 0) {
        cout << 0 << endl;
        continue;
      }
      int64_t cv = nd - vds[i];
      while (cntd <= cv)
        cv -= cntd;
      while (cv < 0)
        cv += cntd;
      // cout << i << " " << nd << " " << cv << " " << cntd << endl;
      cout << solve(cv) + 1 << endl;
    } else {
      int64_t cv = nu + vus[i];
      while (cntu <= cv)
        cv -= cntu;
      while (cv < 0)
        cv += cntu;
      // cout << i << " " << nu << " " << cv << " " << cntu << endl;
      cout << solve(cv) + 1 << endl;
    }
  }

  return 0;
}