ABC451D - Concat Power of 2 を Python で解いた

D - Concat Power of 2
https://atcoder.jp/contests/abc451/tasks/abc451_d

1ヶ月ぶりの記事です。AtCoder のやる気を失ったわけではないのですが記事を書く余裕がなさすぎてサボっています。AWC (AtCoder Weekday Contest) が登場したことにより、記事を書く暇があったら問題をどんどん解きたいという欲求が大きくなってしまいました。まあ参加した回の振り返りも余裕があるときには書きたいと思います。

コンテスト中の動き

12分でCまで終わらせてDを解き始めました。最初は「1, 2, 4, 8 は出てくる、0, 3, 5, 7, 9 は出てこない、6 は 1 桁では使えないが 16 のような形では出てくる」という浅い理解をしていました。だから「1, 2, 4, 8 の 4 進数だけど $ 2^4, 2^8, 2^{12}, ... $ の場合のみ 6 が出てくる」という法則だと思い込んでしまいました。もちろん誤りです。自分の思い込んでいた法則の通りに実装して列挙してみて、入力例 2 と合わないのでようやく誤りに気づきます。確かに 256 みたいに 5 とかも普通に出てくるし、6 だって 116 みたいな形でも現れます。これに気づくまでに 17 分。

次に、桁ごとに個数を計算できるのではないかと考えて 1-3 桁の場合を考えたりしていましたが無理そうだなと思って断念。ここまでで 9分ほど。

Eに浮気しました。あれこれコードを書きながら考えた挙げ句、諦めます。27 分消費しましたが D もわからなかったので仕方ありません。

再び D に帰ってきて、さっき考えていたメモを見返すうちに k-1 桁以下の良い数から k 桁の良い数を作れるのでは？という考えに至りました。もっというと、足して k になるように k-1 以下の数字を複数選んで足し合わせる。ただ、それだと計算量が爆発する（場合の数がとんでもないことになる）からどうかなと思いました。しかし、もっとシンプルに考えられると気づいてそのように実装し、なんとか正解できました。Dに帰ってきてから 12 分、なんとか AC です。コンテスト開始から 77 分もかかってしまいました。レート変化は横ばい。もうずいぶん長い間 1000 - 1050 ぐらいをうろうろしています。

考察

大まかな流れはさっき書いた通りですがもう少し詳しく書きます。入力例 3 の答えが 819,264,512 、およそ $ 8 * 10^8 $ なので $ 10^9 $ に近く、ほぼ最大だと勝手に予想します。その際の入力が 1,099,898 なのでおよそ $ 1.1 * 10^6 $、全部列挙しても間に合うと予想しました。

• 1桁の良い数 ... 1, 2, 4, 8 の 4 つ
• 2桁の良い数 ... べき乗が 16, 32, 64 の 3 つ。1桁の良い数を2つ組み合わせると $ 4\times4=16 $ 個。計 19 個。実際書き出してみると一致する。
• 3桁の良い数 ... 128, 256, 512 の 3つ。128は1桁の良い数の組み合わせで作れる。重複は後で set() で排除すればいいか。1桁の良い数と2桁の良い数をくっつけて3桁の良い数が作れる。逆に2桁+1桁という組み合わせもある。あとは1桁+1桁+1桁で作れる。

……この方針でいくと 4 桁以降どんどん計算量が増えていきますね。9 桁になると 1+1+1+.... とか 3+3+3 とか 2+2+3+2 とかいろんな組み合わせが出てきます。これを全列挙するコードを書けと言われても難しいですし、おそらく計算時間も間に合わないです。

（Eに浮気した後戻ってきて）ここであることに気づきます。k 桁の良い数を作るとき、k-d 桁の良い数と d 桁の良い数（$ 1 \leq d \leq k-1 $）を足し合わせる場合だけ考えればいいのでは？ と。

9桁の場合なら8桁+1桁、7桁+2桁、6桁+3桁、5桁+4桁、1桁+8桁、2桁+7桁、3桁+6桁、4桁+5桁。

例えば 2+2+2+2+1 のような場合、8桁の良い数の中に既に 2+2+2+2 の部分が内包されているのだから改めて数え上げる必要はないということです。確信はないけどやってみることにしました。そしたら見事に AC することができました。

実装

桁ごとに良い数をまとめる変数 D を用意し、最初にべき乗を追加。そのあと 2 桁から 9 桁まで列挙していきます。最後は全ての良い数を一つのリスト E にまとめて、全部いっぺんにソートしています。ちなみに実行時間は 1997ms。一応 AC は AC なので別にギリギリでもいいのですが、桁ごとにソートした方が速かったかもしれませんね。

N = int(input())
D = [set() for _ in range(11)] # digit
for i in range(31):
  num = 2 ** i
  l = len(str(num))
  D[l].add(str(num))

for d in range(2, 10):
  # (d - 1, 1)
  # (d - 2, 2)
  # ..........
  for i in range(1, d //2 + 1):
    for a in D[d-i]:
      for b in D[i]:
        D[d].add(str(a)+str(b))
        D[d].add(str(b)+str(a))

E = []
for i in range(11):
  for j in D[i]:
    E.append(int(j))
E.sort()
print(E[N-1])