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

#include <bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef vector<vector<int>> vvi;
typedef pair<int, int> pii;
typedef long long ll;
const long long MOD = 1000000007;
#define rep(i, n) for (long long i = 0; i < (n); i++)
#define rep2(i, m, n) for (long long i = (m); i < (n); i++)
#define ALL(v) v.begin(), v.end()
#define pb push_back

// 2のa乗をbで割った余り

ll amari(ll a, ll b) {
  ll tmp = 1 % b;

  rep(i, a) { tmp = (tmp * 2) % b; }

  return tmp;
}

int main() {

  ll n;
  cin >> n;
  string x;
  cin >> x;

  ll cnt = 0;

  rep(i, n) {
    if (x.at(i) == '1') {
      cnt++;
    }
  }

  vector<ll> data(200001, 0);

  rep(i, 200001) {
    ll tmp = 0;
    ll a = i;

    while (a != 0) {
      bitset<20> s(a);

      ll cntone = 0;
      rep(j, 20) {
        if (s.test(j)) {
          cntone++;
        }
      }

      a %= cntone;

      tmp++;
    }

    data.at(i) = tmp;
  }

  if (cnt == 1) {

    rep(i, n) {
      if (x.at(i) == '1') {
        cout << 0 << endl;
      }

      else {
        if (x.at(n - 1) == '0' and i != (n - 1)) {
          cout << 1 << endl;
        } else if (x.at(n - 1) == '1' and i == n - 1) {
          cout << 1 << endl;
        } else {
          cout << 2 << endl;
        }
      }
    }
  }

  else if (cnt == 0) {
    rep(i, n) { cout << 1 << endl; }
  }

  else {

    vector<ll> minusamari(n, 0);
    vector<ll> plusamari(n, 0);

    ll tmpminus = 1 % (cnt - 1);
    ll tmpplus = 1 % (cnt + 1);

    minusamari.at(0) = tmpminus;
    plusamari.at(0) = tmpplus;

    rep(i, n - 1) {

      tmpminus = (tmpminus * 2) % (cnt - 1);
      tmpplus = (tmpplus * 2) % (cnt + 1);

      minusamari.at(i + 1) = tmpminus;
      plusamari.at(i + 1) = tmpplus;
    }

    ll minusfirstamari = 0;
    ll plusfirstamari = 0;

    rep(i, n) {
      ll tmp;
      if (x.at(i) == '1') {
        tmp = minusamari.at(n - 1 - i);
        minusfirstamari += tmp;
        minusfirstamari %= (cnt - 1);

        tmp = plusamari.at(n - 1 - i);
        plusfirstamari += tmp;
        plusfirstamari %= (cnt + 1);
      }
    }

    rep(i, n) {
      if (x.at(i) == '1') {

        ll zantei;

        zantei = minusfirstamari - minusamari.at(n - 1 - i);

        if (zantei < 0) {
          zantei += cnt - 1;
        }

        cout << 1 + data.at(zantei) << endl;
        ;

      }

      else {

        ll zantei;
        zantei = plusfirstamari + plusamari.at(n - 1 - i);
        zantei %= (cnt + 1);

        cout << 1 + data.at(zantei) << endl;
      }
    }
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef vector<int> vi;
typedef vector<vector<int>> vvi;
typedef pair<int, int> pii;
typedef long long ll;
const long long MOD = 1000000007;
#define rep(i, n) for (long long i = 0; i < (n); i++)
#define rep2(i, m, n) for (long long i = (m); i < (n); i++)
#define ALL(v) v.begin(), v.end()
#define pb push_back

// 2のa乗をbで割った余り

ll amari(ll a, ll b) {
  ll tmp = 1 % b;

  rep(i, a) { tmp = (tmp * 2) % b; }

  return tmp;
}

int main() {

  ll n;
  cin >> n;
  string x;
  cin >> x;

  ll cnt = 0;

  rep(i, n) {
    if (x.at(i) == '1') {
      cnt++;
    }
  }

  vector<ll> data(200000, 0);

  rep(i, 200000) {
    ll tmp = 0;
    ll a = i;

    while (a != 0) {
      bitset<20> s(a);

      ll cntone = 0;
      rep(j, 20) {
        if (s.test(j)) {
          cntone++;
        }
      }

      a %= cntone;

      tmp++;
    }

    data.at(i) = tmp;
  }

  if (cnt == 1) {

    rep(i, n) {
      if (x.at(i) == '1') {
        cout << 0 << endl;
      }

      else {
        if (x.at(n - 1) == '0' and i != (n - 1)) {
          cout << 1 << endl;
        } else if (x.at(n - 1) == '1' and i == n - 1) {
          cout << 1 << endl;
        } else {
          cout << 2 << endl;
        }
      }
    }
  }

  else if (cnt == 0) {
    rep(i, n) { cout << 1 << endl; }
  }

  else {

    ll minusfirstamari = 0;
    ll plusfirstamari = 0;

    rep(i, n) {
      ll tmp;
      if (x.at(i) == '1') {
        tmp = amari(n - 1 - i, cnt - 1);
        minusfirstamari += tmp;
        minusfirstamari %= (cnt - 1);

        tmp = amari(n - 1 - i, cnt + 1);
        plusfirstamari += tmp;
        plusfirstamari %= (cnt + 1);
      }
    }

    rep(i, n) {
      if (x.at(i) == '1') {

        ll zantei;

        zantei = minusfirstamari - amari(n - 1 - i, cnt - 1);

        if (zantei < 0) {
          zantei += cnt - 1;
        }

        cout << 1 + data.at(zantei) << endl;
        ;

      }

      else {

        ll zantei;
        zantei = plusfirstamari + amari(n - 1 - i, cnt + 1);
        zantei %= (cnt + 1);

        cout << 1 + data.at(zantei) << endl;
      }
    }
  }
}