【全探索】ARC034 B - 方程式をPythonで解いてみた

atcoder.jp

こんにちは。今日は上記の問題を解いてみます。
例の通り、問題は上記のリンクから読んでください。

B問題といってもARC。難易度は高そうです。

思考の過程

なるほど。
まず目につくのは、標準入力のNのオーダーですね。
最大10の18乗ということは、単純にiを1からNまで回すfor文では、TLEとなる問題が存在しそうです。
そこに気をつけながら解いていきましょう。

さて、問題の大まかな方針についてですが、与えられた等式が成り立つか否かのif文で良さそうですね。

その際、各桁の数の和が厄介です。
$x$ によって $f(x)$ の値も変わってしまうので、 $x$ を1からNまで変化させてすべて試したくなりますが、先程のNの最大値が10の18乗という制約があるため、それはACしません。

じゃあ、その厄介な $f(x)$ について詳しく見てみましょう。
この $f(x)$ 、そんなに大きい値にはならなそうじゃないですか？
だって、単なる各桁の和ですよ？ $x=999999$ でさえ、たったの $54$ です。

つまりこの $f(x)$ は、最大でも $17×9=153$ です。*1

したがって、等式が取りうる大体の範囲が分かれば、その範囲を決め打ちして、for文で全探索できそうです。
この大体の範囲は実はもうわかります。
というのも、 $f(x)$ の最大は $153$ だからです。

つまり、 $x$ の値を $N-153$ から $N$ までfor文で全探索すれば良さそうです。

ですが、ここで1つ気をつけなければならないこととして、 $N<153$ のときです。

このとき、for文の始点が負となってしまうので、分けて実装するべきでしょう。
この程度の範囲なら決め打ちする必要はないですしね。

じゃあ実装します。

コード実装

プログラムの流れは以下のようにする。
1. Nの大きさで場合分け(今回は1000以下か否か)
2. Nが1000以下の場合、1からNまで等式が成り立つかを計算
3. Nが1000以上の場合、N-200(153だと中途半端なので)からNまで等式が成り立つかを計算
4. 等式が成り立つ場合、その値をリストに格納

n = input()
N = int(n)
cnt = 0
x_list = []
if N <= 1000:  # Nの大きさで場合分け
    for i in range(N):  # Nが1000以下の場合、1からNまで等式が成り立つかを計算
        sum_ = 0
        tmp = str(i)
        for j in range(len(tmp)):
            sum_ += int(tmp[j])
        if i + sum_ == N:
            cnt += 1
            x_list.append(i)  # 等式が成り立つ場合、その値をリストに格納
else:
    for i in range(N-200, N):  # Nが1000以上の場合、N-200からNまで等式が成り立つかを計算
        sum_ = 0
        tmp = str(i)
        for j in range(len(tmp)):
            sum_ += int(tmp[j])
        if i + sum_ == N:
            cnt += 1
            x_list.append(i)  # 等式が成り立つ場合、その値をリストに格納
print(cnt)
for i in range(len(x_list)):
    print(x_list[i])