けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 161 D - Lunlun Number (400 点)

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

問題へのリンク

問題概要

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

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

制約

  •  1 \le K \le 10^{5}

考えたこと

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


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

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

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

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

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

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

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

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

drken1215.hatenablog.com

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

f:id:drken1215:20200405175042p:plain

まず最初に 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 桁以下のルンルン数を全列挙して、そのうちの  K 番目を求めるという方針だ。

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

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

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

#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 になる)

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

drken1215.hatenablog.com

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


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

そうして  K 回目に取り出された値が答えとなるのだ。さらによくよく考えてみよう。当初は 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

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


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


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


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


さて、この問題は見るからに桁 DP で解けそうだ。一般に「 x 以下の整数であって、〜という条件を満たすものは何個あるか」という問題は桁 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;
}