Problem p02609 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
#define ff first
#define ss second
// #define pb push_back
#define debug(val) cerr << "The value of " << #val << " is = " << val << '\\n';
typedef long long ll;
typedef long double ld;
const ld pi = 4 * atan((ld)1);
const ll mod = 1e9 + 7;
const ll inf = 922337203685477;
const ll nax = 1e6 + 5;

ll n, a[nax], b[nax], one, zero, ans[nax];
string s;

int countSetBits(int x) {
  if (x == 0)
    return 0;
  else
    return 1 + countSetBits(x & (x - 1));
}

ll f(ll x) {
  if (x == 0) {
    return 0ll;
  } else {
    ll y = countSetBits(x);
    x = x % y;
    return 1ll + f(x);
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);

  cin >> n;
  cin >> s;
  reverse(s.begin(), s.end());
  for (ll i = 0; i < n; i++) {
    if (s[i] == '1') {
      one++;
    } else {
      zero++;
    }
  }
  ll A = 0, B = 0;
  ll pa = 1;
  ll pb = 1;
  if (one <= 1) {
    pb = 0;
  }
  for (ll i = 0; i < n; i++) {
    a[i] = pa;
    if (s[i] == '1') {
      A = (A + pa) % (one + 1);
    }
    pa = (pa * 2) % (one + 1);
    if (one <= 1) {
      continue;
    } else {
      b[i] = pb;
      if (s[i] == '1') {
        B = (B + pb) % (one - 1);
      }
      pb = (pb * 2) % (one - 1);
    }
  }
  for (ll i = 0; i < n; i++) {
    if (s[i] == '1') {
      if (one == 1) {
        ans[i] = 0;
      } else {
        ll bb = B - b[i];
        bb = (bb + (one - 1)) % (one - 1);
        ll ret = f(bb);
        ans[i] = ret + 1;
      }
    } else {
      ll aa = A + a[i];
      aa = (aa + one + 1) % (one + 1);
      ll ret = f(aa);
      ans[i] = ret + 1;
    }
  }
  for (ll i = n - 1; i >= 0; i--) {
    cout << ans[i] << '\\n';
  }
}


Error: WA

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define ff first
#define ss second
// #define pb push_back
#define debug(val) cerr << "The value of " << #val << " is = " << val << '\\n';
typedef long long ll;
typedef long double ld;
const ld pi = 4 * atan((ld)1);
const ll mod = 1e9 + 7;
const ll inf = 922337203685477;
const ll nax = 1e6 + 5;

ll n, a[nax], b[nax], one, zero, ans[nax];
string s;

int countSetBits(int x) {
  if (x == 0)
    return 0;
  else
    return 1 + countSetBits(x & (x - 1));
}

ll f(ll x) {
  if (x == 0) {
    return 0ll;
  } else {
    ll y = countSetBits(x);
    x = x % y;
    return 1ll + f(x);
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);

  cin >> n;
  cin >> s;
  reverse(s.begin(), s.end());
  for (ll i = 0; i < n; i++) {
    if (s[i] == '1') {
      one++;
    } else {
      zero++;
    }
  }
  ll A = 0, B = 0;
  ll pa = 1;
  ll pb = 1;
  if (one <= 1) {
    pb = 0;
  }
  for (ll i = 0; i < n; i++) {
    a[i] = pa;
    if (one != 0) {
      b[i] = pb;
    }
    if (s[i] == '1') {
      A = (A + pa) % (one + 1);
      if (one != 0) {
        B = (B + pb) % (one - 1);
      }
    }
    pa = (pa * 2) % (one + 1);
    if (one != 0) {
      pb = (pb * 2) % (one - 1);
    }
  }
  // for(ll i = 0; i < n; i++){
  // if(s[i] == '1'){
  // if(one == 1){
  // ans[i] = 0;
  //}
  // else{
  // ll bb = B - b[i];
  // bb = (bb + (one - 1)) % (one - 1);
  // ll ret = f(bb);
  // ans[i] = ret + 1;
  //}
  //}
  // else{
  // ll aa = A + a[i];
  // aa = (aa + one + 1) % (one + 1);
  // ll ret = f(aa);
  // ans[i] = ret + 1;
  //}
  //}
  for (ll i = n - 1; i >= 0; i--) {
    cout << ans[i] << '\\n';
  }
}