AtCoder ABC 161 D - Lunlun Number (緑色, 400 点)

この手の DFS、一時期全然見なかったけど、最近復活してきた！

問題へのリンク

問題概要

1 以上の整数であって、隣り合う桁の値の絶対値が 1 以下であるような数をルンルン数とよぶ。

番目に小さいルンルン数を求めよ。

制約

• 

考えたこと

さて、「 番目に小さいものを求めよ」という問題では、ザッと次のような方針で解けることが多い！！！

---

1. 考えられる数を全列挙して、ソートして、 番目を取り出す
2. priority_queue を使いながら列挙していく
3. 二分探索する

---

今回は、1 の方針で解いた人が多かったようだ。2 の方針でもできる。そして、3 つめの方針は「 番目を求めよ」という問題では最も多いパターンだと思われる。が、今回の問題ではその方針は大変なことになる。今回の問題では、

• 1 でやると、DFS になる
• 2 でやると、結果的に priority_queue 要らないじゃん、queue でいいじゃん、という流れで BFS になる
• 3 でやると、二分探索 + 桁 DP というえげつないことになる

なお、1 の方針でも、再帰関数を使う書き方の他に、再帰関数を使わない書き方も考えられる。

解法 1-1： DFS で全列挙、再帰関数を使う

サンプル 3 を見てみよう。最大の に対する答えが、3234566667 だとわかる！！！つまり、最大でも 10 桁以下のルンルン数をすべて列挙しておけば問題なさそうだ。

一応真面目に考えると、 桁のルンルン数はだいたい 個くらいになりそう。このことから くらいで十分そう、という予想を立てることはできる。

それでは、10 桁以下のルンルン数を全列挙しよう。こういう場面では再帰関数が大活躍するのだ。類題としては、こんな問題がある。

drken1215.hatenablog.com

場合の数とかの問題で、樹形図というのを描いたことはあるかもしれない。まず樹形図というのを思い浮かべてみよう。今回の問題の場合は下図のような感じ。

まず最初に 1〜9 を列挙しておいて、そこから次々と新しい数値を付け加えていくイメージ。こういう樹形図は「木」なので、DFS や BFS で辿ることができるのだ。ここでは DFS で辿ってみる。

コツは、こういう感じの再帰関数を作るのだ。この再帰の書き方は、色々応用が利くと思う。

// d: 現在のケタ数
// val: 現在の値
// all: ここに格納していく (参照型にしておく)
void rec(int d, long long val, vector<long long> &all) {
  // 格納
  all.push_back(val);

  // 10 桁だったらそれ以上やらずに打ち切り
  if (d == 10) return;

  // 次の桁へ進む処理
  (ここで再帰呼び出しをする)
} 

このフォーマットを意識して、書いてみよう。

• 格納
• 終端条件を満たしたら打ち切り
• 再帰呼び出し

をそれぞれ埋めていくイメージだ。

#include <bits/stdc++.h>
using namespace std;

void rec(int d, long long val, vector<long long> &all) {
    // 格納
    all.push_back(val);

    // 10 桁だったらそれ以上やらずに打ち切り
    if (d == 10) return;

    // 次の桁へ進む処理
    for (int j = -1; j <= 1; ++j) {
        int add = (val % 10) + j;
        if (add >= 0 && add<= 9) rec(d+1, val*10+add, all);
    }
}

int main() {
    int K;
    cin >> K;
    vector<long long> all;
    for (int v = 1; v < 10; ++v) rec(1, v, all);

    // 小さい順に並び替える
    sort(all.begin(), all.end());

    // K 番目
    cout << all[K-1] << endl;
}

解法 1-2： 再帰関数を用いずに全列挙

解法 1-1 と基本的には同じ方針だ。10 桁以下のルンルン数を全列挙して、そのうちの 番目を求めるという方針だ。

それは次のようにすれば、再帰関数を使わずに列挙することもできる。次のような関数を用意するのだ。

• 桁のルンルン数がすべれ列挙された配列 ar を受け取って
• 桁のルンルン数をすべて列挙した配列を求める

この方針で実装してみよう。ソースコードを見てしまうのが早いと思う。

#include <bits/stdc++.h>
using namespace std;

// d 桁のルンルン数の集合から、d+1 桁のルンルン数の集合を求める
vector<long long> calc_next(const vector<long long> &ar) {
    vector<long long> res;
    for (auto val : ar) {
        for (int j = -1; j <= 1; ++j) {
            int add = (val % 10) + j;
            if (add >= 0 && add <= 9) res.push_back(val * 10 + add);
        }
    }
    return res;
}

int main() {
    int K;
    cin >> K;
    vector<long long> all; // 全列挙したもの

    // 一桁の場合
    vector<long long> cur;
    for (int v = 1; v < 10; ++v) cur.push_back(v), all.push_back(v);

    // 10 桁まで
    for (int d = 1; d < 10; ++d) {
        // 次の桁を列挙
        cur = calc_next(cur);

        // all に格納
        for (auto val : cur) all.push_back(val);
    }

    // K 番目
    cout << all[K-1] << endl;
}

解法 2： priority_queue を使う (結局 queue になる)

「 番目に小さいものを求めよ」という問題で、priority_queue が有効に使えるケースは結構多い。具体例として次のような問題例がある。

drken1215.hatenablog.com

こんな風にしていくのだ。

---

• まず (プライオリティ) キューに 1〜9 を push する
• 以下の操作を 回繰り返す
  • キューの先頭を取り出す ( とする)
  • の一の位の値が であるとき、
    • > 0 なら、キューに を push
    • キューに を push
    • < 9 なら、キューに を push

---

そうして 回目に取り出された値が答えとなるのだ。さらによくよく考えてみよう。当初は priority_queue を使うという風に考えていたのだが、上の構成方法をよく眺めると

• キューに push される値は、時系列に沿って単調増加である

ということがわかる。よって、実は priority_queue ではなく、ただの queue で構わない！！！！！！

そしてこの解法は、解法 1-1 でやったような「樹形図上の DFS」を、「樹形図上の BFS」に変更したものと捉えることもできる！

#include <bits/stdc++.h>
using namespace std;

int main() {
    int K;
    cin >> K;
    queue<long long> que;
    for (int i = 1; i < 10; ++i) que.push(i);
    for (int i = 0; i < K-1; ++i) {
        long long s = que.front();
        que.pop();
        for (int j = -1; j <= 1; ++j) {
            long long add = (s % 10) + j;
            if (add >= 0 && add <= 9) que.push(s * 10 + add);
        }
    }
    cout << que.front() << endl;
}

解法 3： 二分探索 + 桁 DP

一般に 番目に小さい値を求めよ、と問われたら、最も多い解法パターンは二分探索だと思う。二分探索をするとこういう風に言い換えることができる

---

番目の値が である
⇔ 「 以下で条件を満たすものが 個ある」という条件を満たす最小の整数 が である

---

この言い換えをすると、結局「 以下の条件を満たすものが何個あるかを求めよ」という問題に落ちるのだ。今回でいえば、

---

以下のルンルン数が何個あるのかを求めよ ( 個以上かどうかを判定せよ)

---

さて、この問題は見るからに桁 DP で解けそうだ。一般に「 以下の整数であって、〜という条件を満たすものは何個あるか」という問題は桁 DP で解くことができることが多い。

drken1215.hatenablog.com

今回はこんな感じでできる。めちゃくちゃ煩雑で辛い。

---

dp[ i ][ j ][ smaller ] := i 桁目まで決めたとき、それを x の上から i 桁目までと比較したときに、

• smaller = 0 のとき、x とちょうど一致する場合
• smaller = 1 のとき、x よりも小さくなっている場合

についての、

• j = 0〜9 のとき、最終桁が j である場合
• j = 10 のとき、i 桁目までがすべて 0 である場合

についての、ルンルン数の個数

---

#include <bits/stdc++.h>
using namespace std;

// x 以下のルンルン数が何個あるか
long long calc(long long x) {
    // x を文字列にする
    stringstream ss;
    ss << x;
    string S = ss.str();
    int N = S.size();

    long long dp[20][11][2] = {0};
    dp[1][10][1] = 1;
    for (int j = 1; j < S[0]-'0'; ++j) dp[1][j][1] = 1;
    dp[1][S[0]-'0'][0] = 1;
    for (int i = 1; i < N; ++i) {
        // from j = 10
        dp[i+1][10][1] += dp[i][10][0] + dp[i][10][1];
        for (int add = 1; add < 10; ++add)
            dp[i+1][add][1] += dp[i][10][0] + dp[i][10][1];

        // from j = 0 ~ 9
        for (int j = 0; j < 10; ++j) {
            for (int k = -1; k <= 1; ++k) {
                int add = j + k;
                if (add < 0 || add > 9) continue;

                dp[i+1][add][1] += dp[i][j][1];
                if (add < S[i]-'0') 
                    dp[i+1][add][1] += dp[i][j][0];
                else if (add == S[i]-'0')
                    dp[i+1][add][0] += dp[i][j][0];
            }
        }
    }
    long long res = 0;
    for (int j = 0; j < 10; ++j) res += dp[N][j][0] + dp[N][j][1];
    return res;
}

int main() {
    int K;
    cin >> K;
    long long left = 0, right = 1LL<<50;
    while (right - left > 1) {
        long long x = (left + right) / 2;
        if (calc(x) >= K) right = x;
        else left = x;
    }
    cout << right << endl;
}