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

#include <bits/stdc++.h>
#ifdef __LOCAL
#define DBG(X) cout << #X << " = " << (X) << endl;
#define SAY(X) cout << (X) << endl;
#else
#define DBG(X)
#define SAY(X)
#endif

#ifdef __LOCAL
#include <filesystem>
namespace fs = std::filesystem;
#endif

using namespace std;

using ll = long long int;
using ull = unsigned long long int;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;

int dx[] = {1, 0, -1, 0, 1, -1, 1, -1};
int dy[] = {0, 1, 0, -1, 1, 1, -1, -1};
const int INT_INF = (int)(2e9);
const ll LL_INF = (ll)(2e18);
static mt19937 _g(time(nullptr));

inline ll randint(ll a, ll b) {
  ll w = (_g() << 31LL) ^ _g();
  return a + w % (b - a + 1);
}
inline void fast_io() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
};
template <typename T, typename S>
inline ostream &operator<<(ostream &os, const pair<T, S> p) {
  cout << "[" << p.first << ";" << p.second << "]";
  return os;
}
template <typename T, typename S>
inline ostream &operator<<(ostream &os, const map<T, S> p) {
  for (auto el : p)
    cout << "[" << el.first << ";" << el.second << "]";
  return os;
}
template <typename T>
inline ostream &operator<<(ostream &os, const vector<T> &v) {
  for (auto el : v)
    cout << el << " ";
  return os;
}
template <typename T> inline vector<T> fetch_vec(int sz) {
  vector<T> ret(sz);
  for (auto &elem : ret)
    cin >> elem;
  return ret;
}

int K;
void input() {
  fast_io();
#ifdef __LOCAL
  fs::path p = __FILE__;
  fs::path input, output;
  input = output = p.parent_path();
  input += string("/input/") + string(p.stem()) + string(".txt");
  output += string("/output/") + string(p.stem()) + string(".txt");
  freopen(input.c_str(), "r", stdin);
  freopen(output.c_str(), "w", stdout);
#endif
  cin >> K;
}

// x以下のるんるん数の個数
// O(|size(x)|)
ll numOfLunlun(ll x) {
  string z = to_string(x);
  int n = z.length();
  ll dp[100][2][10];
  memset(dp, 0, sizeof(dp));

  dp[0][0][(z[0] - '0')] = 1;
  for (int i = 0; i < (z[0] - '0'); i++) {
    dp[0][1][i] = 1;
  }

  for (int i = 0; i < n - 1; i++) {
    for (int k = 0; k < 10; k++) {
      dp[i + 1][1][k] += dp[i][1][k];
      if (k != 0) {
        dp[i + 1][1][k] += 1; // 000000kの追加
        dp[i + 1][1][k] += dp[i][1][k - 1];
        if (k == 1)
          dp[i + 1][1][k]--; // 0000001の重複分削除
      }
      if (k != 9)
        dp[i + 1][1][k] += dp[i][1][k + 1];
      // if(k<(z[i+1]-'0')) dp[i+1][1][k] += dp[i][0][k];
      // if(k<=(z[i+1]-'0') && k!=0) dp[i+1][1][k] += dp[i][0][k-1];
      // if(k+1<(z[i+1]-'0') && k!=9) dp[i+1][1][k] += dp[i][0][k+1];
    }
    if (0 <= (z[i] - '0' - 1) and (z[i] - '0' - 1) < (z[i + 1] - '0'))
      dp[i + 1][1][z[i] - '0' - 1] += dp[i][0][z[i] - '0'];
    if ((z[i] - '0' + 0) < (z[i + 1] - '0'))
      dp[i + 1][1][z[i] - '0' + 0] += dp[i][0][z[i] - '0'];
    if ((z[i] - '0' + 1) < (z[i + 1] - '0'))
      dp[i + 1][1][z[i] - '0' + 1] += dp[i][0][z[i] - '0'];

    if (abs(z[i + 1] - z[i]) <= 1)
      dp[i + 1][0][z[i + 1] - '0'] = dp[i][0][z[i] - '0'];

    // dp[i+1][0][z[i+1]-'0'] += dp[i][0][z[i+1]-'0'];
    // if((z[i+1]-'0'+1)<10) dp[i+1][0][z[i+1]-'0'] += dp[i][0][z[i+1]-'0'+1];
    // if((z[i+1]-'0'-1)>=0) dp[i+1][0][z[i+1]-'0'] += dp[i][0][z[i+1]-'0'-1];
  }
  ll ret = 0;
  for (int k = 0; k < 10; k++) {
    ret += dp[n - 1][0][k] + dp[n - 1][1][k];
  }
  return ret - 1; // -1は000000の削除分
}

int solve(int K = K) {
  ll ng = 0;
  ll ok = 2e18;
  while (ok - ng > 1) {
    ll mid = (ok + ng) / 2;
    if (numOfLunlun(mid) >= K)
      ok = mid;
    else
      ng = mid;
  }
  cout << ok << endl;
  return 0;
}

bool isLunlun(ll x) {
  string z = to_string(x);
  bool res = true;
  for (int i = 0; i < z.length() - 1; i++) {
    res &= (abs(z[i + 1] - z[i]) <= 1);
  }
  return res;
}

ll nextLunlun(ll &x) {
  while (!isLunlun(++x)) {
  }
  return x;
}

int TLEsolve(int k = K) {
  ll res = 1;
  while (--k) {
    nextLunlun(res);
  }
  cout << res << endl;
  return res;
}

int main() {
  input();
  // K = 992;
  solve(K);
  // TLEsolve(K);
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#ifdef __LOCAL
#define DBG(X) cout << #X << " = " << (X) << endl;
#define SAY(X) cout << (X) << endl;
#else
#define DBG(X)
#define SAY(X)
#endif

#ifdef __LOCAL
#include <filesystem>
namespace fs = std::filesystem;
#endif

using namespace std;

using ll = long long int;
using ull = unsigned long long int;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;

int dx[] = {1, 0, -1, 0, 1, -1, 1, -1};
int dy[] = {0, 1, 0, -1, 1, 1, -1, -1};
const int INT_INF = (int)(2e9);
const ll LL_INF = (ll)(2e18);
static mt19937 _g(time(nullptr));

inline ll randint(ll a, ll b) {
  ll w = (_g() << 31LL) ^ _g();
  return a + w % (b - a + 1);
}
inline void fast_io() {
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cout.tie(nullptr);
};
template <typename T, typename S>
inline ostream &operator<<(ostream &os, const pair<T, S> p) {
  cout << "[" << p.first << ";" << p.second << "]";
  return os;
}
template <typename T, typename S>
inline ostream &operator<<(ostream &os, const map<T, S> p) {
  for (auto el : p)
    cout << "[" << el.first << ";" << el.second << "]";
  return os;
}
template <typename T>
inline ostream &operator<<(ostream &os, const vector<T> &v) {
  for (auto el : v)
    cout << el << " ";
  return os;
}
template <typename T> inline vector<T> fetch_vec(int sz) {
  vector<T> ret(sz);
  for (auto &elem : ret)
    cin >> elem;
  return ret;
}

int K;
void input() {
  fast_io();
#ifdef __LOCAL
  fs::path p = __FILE__;
  fs::path input, output;
  input = output = p.parent_path();
  input += string("/input/") + string(p.stem()) + string(".txt");
  output += string("/output/") + string(p.stem()) + string(".txt");
  freopen(input.c_str(), "r", stdin);
  freopen(output.c_str(), "w", stdout);
#endif
  cin >> K;
}

// x以下のるんるん数の個数
// O(|size(x)|)
ll numOfLunlun(ll x) {
  string z = to_string(x);
  int n = z.length();
  ll dp[100][2][10];
  memset(dp, 0, sizeof(dp));

  dp[0][0][(z[0] - '0')] = 1;
  for (int i = 0; i < (z[0] - '0'); i++) {
    dp[0][1][i] = 1;
  }

  for (int i = 0; i < n - 1; i++) {
    for (int k = 0; k < 10; k++) {
      dp[i + 1][1][k] += dp[i][1][k];
      if (k != 0) {
        dp[i + 1][1][k] += 1; // 000000kの追加
        dp[i + 1][1][k] += dp[i][1][k - 1];
        if (k == 1)
          dp[i + 1][1][k]--; // 0000001の重複分削除
      }
      if (k != 9)
        dp[i + 1][1][k] += dp[i][1][k + 1];
      // if(k<(z[i+1]-'0')) dp[i+1][1][k] += dp[i][0][k];
      // if(k<=(z[i+1]-'0') && k!=0) dp[i+1][1][k] += dp[i][0][k-1];
      // if(k+1<(z[i+1]-'0') && k!=9) dp[i+1][1][k] += dp[i][0][k+1];
    }
    if ((z[i] - '0' - 1) < (z[i + 1] - '0'))
      dp[i + 1][1][z[i] - '0' - 1] += dp[i][0][z[i] - '0'];
    if ((z[i] - '0' + 0) < (z[i + 1] - '0'))
      dp[i + 1][1][z[i] - '0' + 0] += dp[i][0][z[i] - '0'];
    if ((z[i] - '0' + 1) < (z[i + 1] - '0'))
      dp[i + 1][1][z[i] - '0' + 1] += dp[i][0][z[i] - '0'];

    if (abs(z[i + 1] - z[i]) <= 1)
      dp[i + 1][0][z[i + 1] - '0'] = dp[i][0][z[i] - '0'];

    // dp[i+1][0][z[i+1]-'0'] += dp[i][0][z[i+1]-'0'];
    // if((z[i+1]-'0'+1)<10) dp[i+1][0][z[i+1]-'0'] += dp[i][0][z[i+1]-'0'+1];
    // if((z[i+1]-'0'-1)>=0) dp[i+1][0][z[i+1]-'0'] += dp[i][0][z[i+1]-'0'-1];
  }
  ll ret = 0;
  for (int k = 0; k < 10; k++) {
    ret += dp[n - 1][0][k] + dp[n - 1][1][k];
  }
  return ret - 1; // -1は000000の削除分
}

int solve() {
  ll ng = 0;
  ll ok = 2e18;
  while (ok - ng > 1) {
    ll mid = (ok + ng) / 2;
    if (numOfLunlun(mid) >= K)
      ok = mid;
    else
      ng = mid;
  }
  cout << ok << endl;
  return 0;
}

bool isLunlun(ll x) {
  string z = to_string(x);
  bool res = true;
  for (int i = 0; i < z.length() - 1; i++) {
    res &= (abs(z[i + 1] - z[i]) <= 1);
  }
  return res;
}

ll nextLunlun(ll &x) {
  while (!isLunlun(++x)) {
  }
  return x;
}

int TLEsolve(int k) {
  ll res = 1;
  while (--k) {
    nextLunlun(res);
  }
  cout << res << endl;
  return res;
}

int main() {
  input();
  // solve();
  TLEsolve(K);
  return 0;
}