早稲田卒のAI研究者のブログ

データサイエンス、機械学習、競技プログラミング、音楽、勉強

【全探索】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])


無事ACでした!


*1 Nが最大のとき、つまりN=10^{18}のときは、f(x)=1