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

#include <bits/stdc++.h>
#define rep(i, n) for (long long i = 0; i < (long long)(n); i++)
#define all(a) a.begin(), a.end()
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<ll, ll> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

const ll MOD = 1e9 + 7;
const ll INF = 1e18;
const int MAX = 510000;
const double pi = acos(-1);
int dx[8] = {1, 0, -1, 0, 1, 1, -1, -1};
int dy[8] = {0, 1, 0, -1, -1, 1, 1, -1};
// int d[13]={0,31,28,31,30,31,30,31,31,30,31,30,31};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  string x;
  cin >> n >> x;
  int num = 0;
  rep(i, n) {
    if (x[i] == '1')
      num++;
  }
  map<ll, ll> mp;
  for (int i = 0; i <= n; i++) {
    int cnt = 0;
    int ni = i;
    while (ni) {
      int bit = __builtin_popcount(ni);
      cnt++;
      ni %= bit;
    }
    mp[i] = cnt;
  }
  vector<ll> p(n), m(n);
  if (num == 1) {
    p[n - 1] = 1;
    ll num1 = 1;
    rep(i, n - 1) {
      num1 *= 2;
      num1 %= (num + 1);
      p[n - 1 - (i + 1)] = num1;
    }
    ll px = 0;
    rep(i, n) {
      if (x[i] == '1') {
        px += p[i];
        px %= (num + 1);
      }
    }
    vector<ll> ans(n);
    rep(i, n) {
      if (x[i] == '1') {
        ans[i] = 0;
      } else {
        ll Px = px;
        Px += p[i];
        Px %= (num + 1);
        while (Px < 0)
          Px += num + 1;
        ans[i] = 1 + mp[Px];
      }
    }
    rep(i, n) cout << ans[i] << endl;
    return 0;
  }
  ll num1 = 1;
  ll num2 = 1;
  p[n - 1] = 1;
  m[n - 1] = 1;
  rep(i, n - 1) {
    num1 *= 2;
    num2 *= 2;
    num1 %= (num + 1);
    num2 %= (num - 1);
    p[n - 1 - (i + 1)] = num1;
    m[n - 1 - (i + 1)] = num2;
  }
  ll px = 0, mx = 0;
  rep(i, n) {
    if (x[i] == '1') {
      px += p[i];
      mx += m[i];
      px %= (num + 1);
      mx %= (num - 1);
    }
  }
  vector<ll> ans(n);
  rep(i, n) {
    if (x[i] == '1') {
      ll Mx = mx;
      Mx -= m[i];
      Mx %= (num - 1);
      while (Mx < 0)
        Mx += num - 1;
      ans[i] = 1 + mp[Mx];
    } else {
      ll Px = px;
      Px += p[i];
      Px %= (num + 1);
      while (Px < 0)
        Px += num + 1;
      ans[i] = 1 + mp[Px];
    }
  }
  rep(i, n) cout << ans[i] << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (long long i = 0; i < (long long)(n); i++)
#define all(a) a.begin(), a.end()
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<ll, ll> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

const ll MOD = 1e9 + 7;
const ll INF = 1e18;
const int MAX = 510000;
const double pi = acos(-1);
int dx[8] = {1, 0, -1, 0, 1, 1, -1, -1};
int dy[8] = {0, 1, 0, -1, -1, 1, 1, -1};
// int d[13]={0,31,28,31,30,31,30,31,31,30,31,30,31};

ll pow(ll x, ll n, ll mo) {
  ll ret = 1;
  while (n > 0) {
    if (n & 1)
      ret = ret * x % mo;
    x = x * x % mo;
    n >> 1;
  }
  return ret;
}
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  string x;
  cin >> n >> x;
  int num = 0;
  rep(i, n) {
    if (x[i] == '1')
      num++;
  }
  map<ll, ll> mp;
  for (int i = 0; i <= n; i++) {
    int cnt = 0;
    int ni = i;
    while (ni) {
      int bit = __builtin_popcount(ni);
      cnt++;
      ni %= bit;
    }
    mp[i] = cnt;
  }
  vector<ll> p(n), m(n);
  if (num == 1) {
    rep(i, n) { p[n - 1 - i] = pow(2, i, num + 1); }
    ll px = 0;
    rep(i, n) {
      if (x[i] == '1') {
        px += p[i];
        px %= (num + 1);
      }
    }
    vector<ll> ans(n);
    rep(i, n) {
      if (x[i] == '1') {
        ans[i] = 0;
      } else {
        ll Px = px;
        Px += p[i];
        Px %= (num + 1);
        while (Px < 0)
          Px += num + 1;
        ans[i] = 1 + mp[Px];
      }
    }
    rep(i, n) cout << ans[i] << endl;
    return 0;
  }
  rep(i, n) {
    p[n - 1 - i] = pow(2, i, num + 1);
    m[n - 1 - i] = pow(2, i, num - 1);
  }
  ll px = 0, mx = 0;
  rep(i, n) {
    if (x[i] == '1') {
      px += p[i];
      mx += m[i];
      px %= (num + 1);
      mx %= (num - 1);
    }
  }
  vector<ll> ans(n);
  rep(i, n) {
    if (x[i] == '1') {
      ll Mx = mx;
      Mx -= m[i];
      Mx %= (num - 1);
      while (Mx < 0)
        Mx += num - 1;
      ans[i] = 1 + mp[Mx];
    } else {
      ll Px = px;
      Px += p[i];
      Px %= (num + 1);
      while (Px < 0)
        Px += num + 1;
      ans[i] = 1 + mp[Px];
    }
  }
  rep(i, n) cout << ans[i] << endl;
  return 0;
}