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

#include <bitset>
#include <iostream>
#include <map>
#include <string>
using namespace std;

std::map<int, long long> mp1;
std::map<int, long long> mp2;
unsigned long ans1 = 0;
unsigned long ans2 = 0;

long long pow(long long n, long long p, long long m) {
  if (m == 0)
    return 0;
  if (p == 0)
    return (1 % m);
  long long ret = 1;
  for (; p > 0; p >>= 1, n = n * n % m) {
    if (p % 2 == 1) {
      ret = ret * n % m;
    }
  }
  return ret;
}

std::map<int, long long> results;

unsigned long f(std::bitset<200000> &bs, int hibit, long long index) {
  long long n = -1;
  long long count = 0;

  // 初回だけmpを使って計算
  // int hibit = bs.count();
  if (bs[index] == 0) {
    hibit++;
    n = ans1 + mp1[index];
  } else {
    hibit--;
    if (hibit <= 0)
      return count;
    n = ans2 - mp2[index];
    if (n < 0)
      n += hibit;
  }
  if (n >= hibit) {
    n = n % hibit;
  }
  count++;

  // 計算量削減のため初回計算後のnとそれに対する計算回数を記憶しておく
  std::map<int, long long>::iterator itr = results.find(n);
  if (itr != results.end()) {
    return itr->second;
  }
  int tmp = n;
  while (n != 0) {
    std::bitset<200000> tmpbs(n);
    n = n % tmpbs.count();
    count++;
  }
  results.emplace(tmp, count);

  // nが0になるまでの操作回数を返す
  return count;
}

int main() {
  long long N;
  std::cin >> N;
  std::string X;
  std::cin >> X;

  // 前準備
  std::bitset<200000> bs(X);
  const int count = bs.count();

  // 制約から1<=Nなので、count=0のケースは考慮不要
  for (long long i = N - 1; i >= 0; i--) {
    mp1[i] = pow(2, i, count + 1);
    mp2[i] = pow(2, i, count - 1);
    if (bs[i] == 1) {
      ans1 += mp1[i];
      ans2 += mp2[i];
    }
  }

  // ここから本処理
  for (long long i = N - 1; i >= 0; i--) {
    unsigned long result = f(bs, count, i);
    std::cout << result << std::endl;
  }
}

Error: TLE

Fixed submission:

#include <bitset>
#include <iostream>
#include <map>
#include <string>
using namespace std;

std::map<int, long long> mp1;
std::map<int, long long> mp2;
unsigned long ans1 = 0;
unsigned long ans2 = 0;

long long pow(long long n, long long p, long long m) {
  if (m == 0)
    return 0;
  if (p == 0)
    return (1 % m);
  long long ret = 1;
  for (; p > 0; p >>= 1, n = n * n % m) {
    if (p % 2 == 1) {
      ret = ret * n % m;
    }
  }
  return ret;
}

std::map<int, long long> results;

unsigned long f(std::bitset<200000> &bs, long long index) {
  long long n = -1;
  long long count = 0;

  // 初回だけmpを使って計算
  int hibit = bs.count();
  if (bs[index] == 0) {
    hibit++;
    n = ans1 + mp1[index];
  } else {
    hibit--;
    if (hibit <= 0)
      return count;
    n = ans2 - mp2[index];
    if (n < 0)
      n += hibit;
  }
  if (n >= hibit) {
    n = n % hibit;
  }
  count++;

  // 計算量削減のため初回計算後のnとそれに対する計算回数を記憶しておく
  std::map<int, long long>::iterator itr = results.find(n);
  if (itr != results.end()) {
    return itr->second;
  }
  int tmp = n;
  while (n != 0) {
    std::bitset<200000> tmpbs(n);
    n = n % tmpbs.count();
    count++;
  }
  results.emplace(tmp, count);

  // nが0になるまでの操作回数を返す
  return count;
}

int main() {
  long long N;
  std::cin >> N;
  std::string X;
  std::cin >> X;

  // 前準備
  std::bitset<200000> bs(X);
  const int count = bs.count();

  // 制約から1<=Nなので、count=0のケースは考慮不要
  for (long long i = N - 1; i >= 0; i--) {
    mp1[i] = pow(2, i, count + 1);
    mp2[i] = pow(2, i, count - 1);
    if (bs[i] == 1) {
      ans1 += mp1[i];
      ans2 += mp2[i];
    }
  }

  // ここから本処理
  for (long long i = N - 1; i >= 0; i--) {
    unsigned long result = f(bs, i);
    std::cout << result << std::endl;
  }
}