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

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cctype>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <list>
#include <numeric>
#include <queue>
#include <regex>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <vector>

#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define sz(x) ((int)(x).size())
#define all(x) (x).begin(), (x).end()

using ll = long long;

ll gcd(ll a, ll b) {
  if (b == 0)
    return a;
  else
    return gcd(b, a % b);
}
ll lcm(ll a, ll b) { return a * b / gcd(a, b); }

int f(int v) {
  if (v == 0)
    return 0;
  int p = std::bitset<32>(v).count();
  return f(v % p) + 1;
}

int main(void) {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  int n;
  std::cin >> n;
  std::string s;
  std::cin >> s;

  // count p
  int p = 0;
  rep(i, sz(s)) if (s[i] == '1') p++;

  // p+1
  int a = 0;
  if (p != n) {
    int x = 1;
    for (int i = sz(s) - 1; i >= 0; i--) {
      if (s[i] == '1')
        a = (a + x) % (p + 1);
      x = (x * 2) % (p + 1);
    }
  }

  // p-1
  int b = 0;
  if (p >= 2) {
    int y = 1;
    for (int i = sz(s) - 1; i >= 0; i--) {
      if (s[i] == '1')
        b = (b + y) % (p - 1);
      y = (y * 2) % (p - 1);
    }
  }
  // std::cout << "\\%p+1 : " << a << std::endl;
  // std::cout << "\\%p-1 : " << b << std::endl;

  // p+1 : a, p-1 : b
  std::vector<int> ans(n);
  int t1 = 1, t2 = 1;
  for (int i = sz(s) - 1; i >= 0; i--) {
    if (s[i] == '1') {
      if (p == 1)
        ans[i] = 0;
      else {
        int next = (b - t1 + (p - 1)) % (p - 1);
        // std::cout << i << " : " << next << std::endl;
        ans[i] = f(next) + 1;
      }
    } else {
      int next = (a + t2) % (p + 1);
      // std::cout << i << " : " << next << std::endl;
      ans[i] = f(next) + 1;
    }
    if (p >= 2)
      t1 = (t1 * 2) % (p - 1);
    t2 = (t2 * 2) % (p + 1);
  }

  for (auto &&m : ans)
    std::cout << m << std::endl;

  return 0;
}


Error: TLE

Fixed submission:


#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cctype>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <list>
#include <numeric>
#include <queue>
#include <regex>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <vector>
// std::numeric_limits<int>::max(), min()
// std::setprecision(X)
// std::islower(), std::isupper(), std::tolower, std::toupper
// std::accumulate(m.begin(), m.end(), 0LL);
// std::sort(begin(),end()) // 0, 1, 2, 3 ,4 (default, less);
// std::priority_queue<int, vector<int>, less<int>> que; // 4, 3, 2, 1 (defualt,
// less) : if ( std::regex_match(S, std::regex("(dream|dreamer|erase|eraser)*"))
// ) std::cout << "YES" << std::endl;

#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define sz(x) ((int)(x).size())
#define all(x) (x).begin(), (x).end()

using ll = long long;

ll gcd(ll a, ll b) {
  if (b == 0)
    return a;
  else
    return gcd(b, a % b);
}
ll lcm(ll a, ll b) { return a * b / gcd(a, b); }

int f(std::string s, int k) {
  // std::cout << "trial : " << k << " , " << s << std::endl;
  int p = 0;
  rep(i, sz(s)) if (s[i] == '1') p++;
  if (p == 0)
    return k;

  int rtn = 0;
  int x = 1;
  for (int i = sz(s) - 1; i >= 0; i--) {
    if (s[i] == '1')
      rtn = (rtn + x) % p;
    x = (x * 2) % p;
  }
  std::bitset<32> b(rtn);
  return f(b.to_string(), k + 1);
}

int main(void) {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  int n;
  std::cin >> n;
  std::string s;
  std::cin >> s;
  std::vector<int> ans(n);
  rep(i, n) {
    std::string v = s;
    if (v[i] == '0')
      v[i] = '1';
    else
      v[i] = '0';
    ans[i] = f(v, 0);
  }

  for (auto &&t : ans)
    std::cout << t << std::endl;
  return 0;
}