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

// 数学をしたほうが実装は楽だと思うけど、さきほどやってみたらバグ死したので、桁DP
// + 2分探索します。
#include <algorithm>
#include <iostream>
#include <vector>
#define int long long
using namespace std;

// a[i] … 10進数における下位i(>=0)桁目.
vector<int> add(vector<int> a, vector<int> b, int n) {
  for (int i = 0; i < n; i++) {
    a[i] += b[i];
    if (i + 1 < n) {
      a[i + 1] += a[i] / 10;
      a[i] %= 10;
    }
  }
  return a;
}

vector<int> sub(vector<int> a, vector<int> b, int n) {
  for (int i = 0; i < n; i++) {
    a[i] -= b[i];
    if (i + 1 < n && a[i] < 0) {
      int kariru = ((-a[i]) + 9) / 10;
      a[i] += 10 * kariru;
      a[i + 1] -= kariru;
    }
  }
  return a;
}

// 2で割る
vector<int> half(vector<int> a, int n) {
  vector<int> ret(n, 0);
  int val = 0;
  for (int i = n - 1; i >= 0; i--) {
    val *= 10;
    val += a[i];
    ret[i] = val / 2;
    val -= ret[i] * 2;
  }
  return ret;
}

int bitCnt[1024];
void initBitCnt() {
  for (int i = 0; i < 1024; i++) {
    bitCnt[i] = 0;
    for (int j = 0; j < 10; j++) {
      if ((i >> j) & 1)
        bitCnt[i]++;
    }
  }
}

// 桁DPで, a以下の短歌数を数える.
int INF = 1e+18;
int dp[101][1024][2]; // dp[残り桁数][今までに使った数の集合(leading
                      // 0はカウントしない)][aと一致しているか]

int dfs(int ketasu, int used, bool flag, vector<int> &a) {
  if (bitCnt[used] > 2)
    return 0;
  if (ketasu == 0)
    return (bitCnt[used] == 2);
  if (dp[ketasu][used][flag] != -1)
    return dp[ketasu][used][flag];

  int ret = 0;
  for (int i = 0; i < 10; i++) { // 次の桁をiにする
    if (flag && a[ketasu - 1] < i)
      break;

    int nUsed;
    if (used == 0 && i == 0)
      nUsed = 0;
    else
      nUsed = (used | (1 << i));
    ret += dfs(ketasu - 1, nUsed, flag && a[ketasu - 1] == i, a);
  }
  return dp[ketasu][used][flag] = ret;
}

// メモの初期化
void initDp() {
  for (int i = 0; i < 101; i++)
    for (int j = 0; j < 1024; j++)
      for (int k = 0; k < 2; k++)
        dp[i][j][k] = -1;
}

// 2分探索の終了条件
bool isContinue(vector<int> st, vector<int> ed, int n) {
  vector<int> res = sub(ed, st, n);

  if (res[0] >= 2)
    return true;
  for (int i = 1; i < n; i++) {
    if (res[i] != 0)
      return true;
  }
  return false;
}

int n;

void solve() {
  initBitCnt();

  int i, cnt = 0;
  for (i = 1;; i++) {
    cnt += 81 * ((1LL << (i - 1)) - 1);
    if (cnt >= n)
      break;
  }
  int ketasu = i;

  vector<int> st(ketasu), ed(ketasu), mid(ketasu);
  for (i = 0; i < ketasu; i++)
    st[i] = 0;
  for (i = 0; i < ketasu; i++)
    ed[i] = 9;

  // 2分探索. xxxooo, (st, ed]
  while (isContinue(st, ed, ketasu)) {
    mid = add(st, ed, ketasu);
    mid = half(mid, ketasu);
    initDp();
    int res = dfs(ketasu, 0, true, mid);
    if (res >= n)
      ed = mid;
    else
      st = mid;
  }

  for (i = ketasu - 1; i >= 0; i--) {
    if (ed[i] != 0)
      break;
  }
  for (; i >= 0; i--) {
    cout << ed[i];
  }
  cout << endl;
}

signed main() {
  while (cin >> n) {
    if (!n)
      break;
    solve();
  }
  return 0;
}


Error: TLE

Fixed submission:

// 数学をしたほうが実装は楽だと思うけど、さきほどやってみたらバグ死したので、桁DP
// + 2分探索します。
#include <algorithm>
#include <iostream>
#include <vector>
#define int long long
using namespace std;

// a[i] … 10進数における下位i(>=0)桁目.
vector<int> add(vector<int> a, vector<int> b, int n) {
  for (int i = 0; i < n; i++) {
    a[i] += b[i];
    if (i + 1 < n) {
      a[i + 1] += a[i] / 10;
      a[i] %= 10;
    }
  }
  return a;
}

vector<int> sub(vector<int> a, vector<int> b, int n) {
  for (int i = 0; i < n; i++) {
    a[i] -= b[i];
    if (i + 1 < n && a[i] < 0) {
      int kariru = ((-a[i]) + 9) / 10;
      a[i] += 10 * kariru;
      a[i + 1] -= kariru;
    }
  }
  return a;
}

// 2で割る
vector<int> half(vector<int> a, int n) {
  vector<int> ret(n, 0);
  int val = 0;
  for (int i = n - 1; i >= 0; i--) {
    val *= 10;
    val += a[i];
    ret[i] = val / 2;
    val -= ret[i] * 2;
  }
  return ret;
}

int bitCnt[1024];
void initBitCnt() {
  for (int i = 0; i < 1024; i++) {
    bitCnt[i] = 0;
    for (int j = 0; j < 10; j++) {
      if ((i >> j) & 1)
        bitCnt[i]++;
    }
  }
}

// 桁DPで, a以下の短歌数を数える.
int INF = 1e+18;
int dp[101][1024][2]; // dp[残り桁数][今までに使った数の集合(leading
                      // 0はカウントしない)][aと一致しているか]

int dfs(int ketasu, int used, bool flag, vector<int> &a) {
  if (ketasu == 0)
    return (bitCnt[used] == 2);
  if (dp[ketasu][used][flag] != -1)
    return dp[ketasu][used][flag];

  int ret = 0;
  for (int i = 0; i < 10; i++) { // 次の桁をiにする
    if (flag && a[ketasu - 1] < i)
      break;

    int nUsed;
    if (used == 0 && i == 0)
      nUsed = 0;
    else
      nUsed = (used | (1 << i));
    ret += dfs(ketasu - 1, nUsed, flag && a[ketasu - 1] == i, a);
  }
  return dp[ketasu][used][flag] = ret;
}

// メモの初期化
void initDp() {
  for (int i = 0; i < 101; i++)
    for (int j = 0; j < 1024; j++)
      for (int k = 0; k < 2; k++)
        dp[i][j][k] = -1;
}

// 2分探索の終了条件
bool isContinue(vector<int> st, vector<int> ed, int n) {
  vector<int> res = sub(ed, st, n);

  if (res[0] >= 2)
    return true;
  for (int i = 1; i < n; i++) {
    if (res[i] != 0)
      return true;
  }
  return false;
}

int n;

void solve() {
  initBitCnt();

  int i, cnt = 0;
  for (i = 1;; i++) {
    cnt += 81 * ((1LL << (i - 1)) - 1);
    if (cnt >= n)
      break;
  }
  int ketasu = i;

  vector<int> st(ketasu), ed(ketasu), mid(ketasu);
  for (i = 0; i < ketasu; i++)
    st[i] = 0;
  for (i = 0; i < ketasu; i++)
    ed[i] = 9;

  // 2分探索. xxxooo, (st, ed]
  while (isContinue(st, ed, ketasu)) {
    mid = add(st, ed, ketasu);
    mid = half(mid, ketasu);
    initDp();
    int res = dfs(ketasu, 0, true, mid);
    if (res >= n)
      ed = mid;
    else
      st = mid;
  }

  for (i = ketasu - 1; i >= 0; i--) {
    if (ed[i] != 0)
      break;
  }
  for (; i >= 0; i--) {
    cout << ed[i];
  }
  cout << endl;
}

signed main() {
  while (cin >> n) {
    if (!n)
      break;
    solve();
  }
  return 0;
}