ABC433D - 183183 を Python で解いた

ABC433 D - 183183
https://atcoder.jp/contests/abc433/tasks/abc433_d

見覚えのあるタイトルと問題文。多くの人が先月解いたばかりのABC428-Dを思い出したことでしょう。

ABC428 D - 183184
https://atcoder.jp/contests/abc428/tasks/abc428_d

記事にしたぐらいなので僕も当然覚えています。

ABC428-D を Python で解いた
https://qiita.com/omakasessan/items/06b5d00e76999e88fac2

コンテスト中の動き

Aがいきなり数学問題で動揺し5分消費。B, Cは簡単だったので7分ほどでAC。Dは26分かけてなんとか解ききることができました。ノーペナで4完でき、しかも久々の水perf。大満足の結果でした。欲を言えば1時間あったんだからE問題は解きたかったですが……実装はおろか、方針も曖昧なまま終わってしまいました。

考察

とりあえずABC428-Dと同じように数式にしてみました。

$$
f(i, j) = A_i \times 10^d + A_j
$$

$(i, j)$ を全探索すると当然 TLE です。この手の数え上げ問題でよくあるのは、「何が何個あるか」という形にして探索ではなく計算で（コンビネーションを使うことも多いですね）答えを求める解法です。その辺から探ってみます。

1. $ A_i $ 側の工夫として、$ A_i $から $ A_i \times 10^{10} $ までが現れるのでそれぞれについて $M$ で割った余りを求めておくことを思いつきます。

2. $ A_j $ 側の工夫として、桁数と $M$ で割った余りの組を作り、それぞれの個数をまとめておくことを思いつきます。が、$M$ は最大で $ 10^9 $にもなるので余りがばらけてしまうとうまく少ない数にまとまってくれない気がします。この方針はなんか違う気がします。

ここで方針を閃きました。ある $ A_j $ に対し、それに対応する $ A_i $ を求めることを考えます。つまり $A_j$ を固定したとき、それに対応する $ A_i \times 10^d $ の個数を数えます。このとき $d$ の値はわかっているわけですから、例えば $A_j$ を $M$ で割った余りが $x$ の場合、$M$ で割った余りが $ M-x $になるような $ A_i \times 10^d $ を探して足せば和が $M$ の倍数になりますね。ここで上記 1. で考えた工夫（前計算）が活きてきます。

試してみる

コンテスト中、念のため入力例2を使って検証してみました。時間はかかりますが確実に正解するために慎重にやりました。

前計算

最上段は配列 A です。下の段はそれぞれ $d$ の値ごとに各 $A_i$ を $M$（この例では$7$） で割った余りです。

		2	8	16	183
d = 1	× 10	6	3	6	3
d = 2	× 100	4	2	4	2
d = 3	× 1000	5	6	5	6

この表を見ながら、各 $A_j$ について対応する $A_i$ を調べてみましょう。$A_j$の桁数を$d$, $M$で割った余りを$m$とします。

• $A_j = 2 $ ... d=1, m=2 → 表の d=1 の行を見て、余り 5 のものを探す。0個。
• $A_j = 8 $ ... d=1, m=1 → 表の d=1 の行を見て、余り 6 のものを探す。2個。
• $A_j = 16 $ ... d=2, m=2 → 表の d=2 の行を見て、余り 5 のものを探す。0個。
• $A_j = 183 $ ... d=3, m=1 → 表の d=3 の行を見て、余り 6 のものを探す。2個。

というわけで計4個の組み合わせが条件を満たすことになり、これは出力例2と合致します。

余りの加算について

ところで、この通りに実装したら入力例3が合いませんでした。少し悩みましたが、余り 0 のパターンの処理が特殊なことに気づきました。この手の問題にありがちですね。M = 7 のとき、

$A_j$ を$M$で割った余り	0	1	2	3	4	5	6
$ A_i \times 10^d $を$M$で割った余り	0	6	5	4	3	2	1

そう、0のときは同じく0が対応しますね。ちなみに$A_j$ を$M$で割った余りと$ A_i \times 10^d $を$M$で割った余りの和は 0 以上 $ 2 \times (M - 1) $ 以下なので、和が 0 もしくは $M$ になる場合だけを考慮すればいいです。

実装

N, M = map(int, input().split())
A = list(map(int, input().split()))
# A に 10**d を掛けてMで割った余りを配列 DA に格納する。
DA = [[0 for _ in range(N)] for _ in range(11)] 
for d in range(11):
  for i in range(N):
    DA[d][i] = (A[i] * 10**d) % M
# DAを見て、桁数ごとに余りいくつのものが何個あるかを数えてまとめる。
# 長さ M の配列にすると TLE するので dict 型を使う。
# dict に入る要素数は最大でも N 個に収まる。
countDA = [dict() for _ in range(11)]  
for d in range(11):
  for i in range(N):
    mod = DA[d][i]
    if mod not in countDA[d]:
      countDA[d][mod] = 0
    countDA[d][mod] += 1

ans = 0
for j in range(N):
  d = len(str(A[j]))
  mod = A[j] % M
  # A[i] * 10**d の表を見て、A[j] に対応するものが何個あるかを調べて加算。
  if mod == 0:
    if 0 in countDA[d]:
      ans += countDA[d][0]
  else:
    if M - mod in countDA[d]:
      ans += countDA[d][M - mod]
print(ans)

ABC428-Dの経験のおかげですんなりと思いつけました。記事を書いて復習したおかげです。