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

#include <algorithm>
#include <cstring>
#include <iostream>
#include <string>

using namespace std;

#define LL long long

const int MAXN = 200000 + 5;

int calPopCount(int N, string S) {
  int ret = 0;
  for (int i = 0; i < N; i++) {
    if (S[i] == '1') {
      ret++;
    }
  }
  return ret;
}

void calMod(int N, string S, int cnt, int *c) {

  if (cnt > 0) {
    c[N - 1] = 1 % cnt;
    for (int i = N - 2; i >= 0; i--) {
      c[i] = (c[i + 1] << 1) % cnt;
    }
  } else {
    for (int i = 0; i < N; i++) {
      c[i] = 0;
    }
  }
}

int f(int m) {
  int cnt = 0;
  for (int i = 0; (1 << i) <= m; i++) {
    if ((1 << i) & m) {
      cnt++;
    }
  }
  return m % cnt;
}

void print(int m) {
  int ans = 1;
  while (m > 0) {
    m = f(m);
    ans++;
  }
  cout << ans << endl;
}

void solve(int N, string S, int cnt, int *a, int *b) {

  int m = 0;
  int ret = 0;
  int aa, bb;
  aa = bb = 0;
  for (int i = 0; i < N; i++) {
    if (S[i] == '1') {
      if (cnt > 1) {
        aa = (aa + a[i]) % (cnt - 1);
      } else {
        aa = 0;
      }

      bb = (bb + b[i]) % (cnt + 1);
    }
  }
  for (int i = 0; i < N; i++) {
    if (S[i] == '0') {
      S[i] = '1';
      int m = (bb + b[i]) % (cnt + 1);

      print(m);
      S[i] = '0';
    } else {
      S[i] = '0';
      int m = 0;
      if (cnt > 1) {
        m = (aa + cnt - 1 - a[i]) % (cnt - 1);
        print(m);
      } else {
        cout << "0" << endl;
      }
      S[i] = '1';
    }
  }
}

int main() {
  int N, a[MAXN], b[MAXN];
  string S;
  cin >> N >> S;

  int cnt = calPopCount(N, S);
  calMod(N, S, cnt - 1, a);
  calMod(N, S, cnt + 1, b);

  solve(N, S, cnt, a, b);

  system("pause");
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstring>
#include <iostream>
#include <string>

using namespace std;

#define LL long long

const int MAXN = 200000 + 5;

int calPopCount(int N, string S) {
  int ret = 0;
  for (int i = 0; i < N; i++) {
    if (S[i] == '1') {
      ret++;
    }
  }
  return ret;
}

void calMod(int N, string S, int cnt, int *c) {

  if (cnt > 0) {
    c[N - 1] = 1 % cnt;
    for (int i = N - 2; i >= 0; i--) {
      c[i] = (c[i + 1] << 1) % cnt;
    }
  } else {
    for (int i = 0; i < N; i++) {
      c[i] = 0;
    }
  }
}

int f(int m) {
  int cnt = 0;
  for (int i = 0; (1 << i) <= m; i++) {
    if ((1 << i) & m) {
      cnt++;
    }
  }
  return m % cnt;
}

int solve(int N, string S, int cnt, int *c) {

  int m = 0;
  int ret = 0;
  // cout << "cnt = " << cnt << endl;
  for (int i = 0; i < N; i++) {
    if (S[i] == '1') {
      m = (c[i] + m) % cnt;
      ret = 1;
      // cout << "m = " << m << endl;
    }
  }

  while (m > 0) {
    m = f(m);
    ret++;
  }
  return ret;
}

int main() {
  int N, a[MAXN], b[MAXN];
  string S;
  cin >> N >> S;

  int cnt = calPopCount(N, S);
  calMod(N, S, cnt - 1, a);
  calMod(N, S, cnt + 1, b);

  for (int i = 0; i < N; i++) {
    if (S[i] == '0') {
      S[i] = '1';
      cout << solve(N, S, cnt + 1, b) << endl;

      S[i] = '0';
    } else {
      S[i] = '0';
      cout << solve(N, S, cnt - 1, a) << endl;
      S[i] = '1';
    }
  }

  system("pause");
  return 0;
}