【Python3】 AtCoder Beginner Contest 161 D – Lunlun Number

更新日 2022-05-25
投稿日 2020-04-04
著者サトゥー
カテゴリー dev

正の整数 $X$ が以下の条件を満たすとき、 $X$ はルンルン数であると言います。
・ $X$ を (leading zeroなしで) 十進数表記した際に、隣り合うどの $2$ つの桁の値についても、差の絶対値が $1$ 以下
例えば、 $1234$ , $1$ , $334$ などはルンルン数ですが、 $31415$ , $119$ , $13579$ などはルンルン数ではありません。
正の整数 $K$ が与えられます。小さい方から $K$ 番目のルンルン数を求めてください。
https://atcoder.jp/contests/abc161/tasks/abc161_d

隣り合う 2 数の差が 1 以下になるような自然数のうち、小さい方から K 番目のものを求める問題です。

サトゥー

とりあえず小さいものからいくつか書き出してみます。例示は理解の試金石、ですよね！

# 1 桁
1, 2, 3, 4, 5, 6, 7, 8, 9

# 2 桁
10, 11, 12, 21, 22, 23, ..., 87, 88, 89, 98, 99

# 3 桁
100, 101, 110, 111, 112, ..., 987, 988, 989, 998, 999

K が大きいので、規則性のようなものを見いだせると良さそうです。そこで桁数に着目して、 $k$ 桁のルンルン数から $k+1$ 桁のルンルン数を生成できないかを考えてみます。

具体例として、ひとまず $k = 1 \to 2$ を考えましょう。

やりたいこと

$k=1$ のルンルン数から $k = 2$ のルンルン数を生成したい

$k = 1$ のルンルン数の後ろに条件を満たすように数字を付けてあげて生成することにします。すると、

1 -> 10, 11, 12
2 -> 21, 22, 23
3 -> 32, 33, 34
4 -> 43, 44, 45
5 -> 54, 55, 56
6 -> 65, 66, 67
7 -> 76, 77, 78
8 -> 87, 88, 89
9 -> 98, 99

のように、 1 ~ 8 からは似たような生成パターンがみられ、 9 からは少し異なる生成パターンがみられることがわかります。

サトゥー

ちなみになんでうしろに数字を付けたかと言うと、前に数字をつけていくと 0 の処理がめんどうだからです

これを一般化して考えます。隣り合う 2 数についての条件ですから、一般の $k \to k + 1$ の場合にも一番うしろの桁（ 1 の位）の値によって生成される数が決まりますから、次のように生成パターンを記述できます。

# k の一番うしろの桁 -> k+1 の一番うしろの 2 桁
0 -> 00, 01
1 -> 10, 11, 12
2 -> 21, 22, 23
3 -> 32, 33, 34
4 -> 43, 44, 45
5 -> 54, 55, 56
6 -> 65, 66, 67
7 -> 76, 77, 78
8 -> 87, 88, 89
9 -> 98, 99

コードに起こすと、次のような関数 next_num()によってまとめられそうです。

def next_num(n):
    if n == 0:
        return [0, 1]
    elif n == 9:
        return [8, 9]
    else:
        return [n-1, n, n+1]

はい、これで生成ルールを決められたので、これによって次の桁のルンルン数を生成しながら、生成した数を小さい順に取り出せば良さそうですね。生成したものを queue に投げ入れておいて、そこから最小値を K 回取り出せばよいので、 heapq を用いて処理をしてみます。

以上、コードはこんな感じです。

from heapq import heapify, heappop, heappush

# 次の桁の生成ルール ex. 1 -> 10, 11, 12
def next_num(n):
    if n == 0:
        return [0, 1]
    elif n == 9:
        return [8, 9]
    else:
        return [n-1, n, n+1]

K = int(input())
queue = [1, 2, 3, 4, 5, 6, 7, 8, 9]  # とりあえず 1 桁の Lunlun Number を初期値にとる
heapify(queue)

for _ in range(K-1):
    now = heappop(queue)  # 最小値を取り出す
    # 取り出した値を使って次の桁の値を生成し queue に投げる
    for next_i in next_num(now%10):
        heappush(queue, now*10 + next_i)

print(queue[0])  # K 回目に取り出す最小値