AtCoder Beginner Contest 161 ABC161 D - Lunlun Number

公式想定解と異なる解き方をしたので共有します。

問題概要

問題へのリンク

正の整数が以下の条件を満たすとき、はルンルン数であるといいます。

• Xを(leading zeroなしで)十進数表記した際に、隣り合うどの2つの桁の値についても、差の絶対値が1以下

小さい方から番目のルンルン数を求めてください。

解説

方針

ある数が番目のルンルン数である条件は、以下のように言い換えられます。

• それ以下にルンルン数が個存在する整数の中で最小である

ある数以下のルンルン数は桁DPで求めることができる（定義から、一つ上の桁のみに依存して下の桁が決まる）ので、この境界を二分探索することができます。

時間計算量はになります。

DPの定義、初期化、遷移

定義

[i][smaller][leadignzeros][lastdigit]上からi桁目までで未満が確定しているか否か(smaller), 0以外を取っているか(leadeingzeros),　最後の桁がlastdigitであるときの場合の数と定義し、DPをします。

初期化

[0][0][0][0]

遷移

実験してみるとわかりますが、まずDPの状態として持ちたいのは最後の桁がいくつであるかというものです。これがわからないと次の桁に使っていい数字がわかりません。またもう一つ重要な情報として、「一度でも0以外を使っているか」という情報があります。これが必要なのは、leadingzeroは含まないと名言されていることに加え、未満フラグが立っている状態で0以外を初めて使う場合には、上の桁に関係なく、0~9何を使ってもよいという性質があるからです。

この状態が多いDPをうまく処理してあげると正確に数え上げることができます。

実装

提出へのリンク

#include <bits/stdc++.h>
using namespace std;
#define rep(i,N) for(int i=0;i<int(N);++i)
#define all(a) (a).begin(),(a).end()
typedef long long ll;

int main() {

    ll K;
    cin >> K;
    // Xを桁ごとのvectorに変換
    auto ll2vec = [&](ll X){
        vector<int> res;
        ll tmp = X;
        while(tmp > 0){
            res.push_back(tmp%10);
            tmp/=10;
        }
        reverse(all(res));
        return res;
    };
    auto check = [&](ll X){
        vector<int> vx = ll2vec(X);
        //桁数
        int N = vx.size();
        // dp[i][smaller][0以外を一度でも取ったか][最後の数字]
        ll dp[N+1][2][2][10];
        rep(i,N+1)rep(j,2)rep(k,2)rep(l,10)
            dp[i][j][k][l] = 0;
        dp[0][0][0][0] = 1;
        //dpの遷移
        for(int i = 0; i < N; i++)rep(j,2)rep(k,2)rep(l,10){
            int min_nxt = max(0, l-1);
            int max_nxt = min(l+1, (j ? 9:vx[i]));
            if(!k)max_nxt = j ? 9:vx[i];
            for(int nxt = min_nxt; nxt <= max_nxt; nxt++){
                dp[i+1][j || nxt < vx[i]][k || nxt][nxt] +=  dp[i][j][k][l];
            }
        }
        ll cnt = 0;
        rep(j,2)rep(l,10)cnt += dp[N][j][1][l];
        return cnt >= K;
    };

    ll ng = 0;
    ll ok = 3234566667;
    while(ok - ng > 1){
        ll mid = (ok + ng) / 2;
        if(check(mid))ok = mid;
        else ng = mid;
    }
    cout << ok << endl;
}