ABC152 D - Handstand 2 - あっとのTECH LOG

問題原文

問題要旨

$N$ 以下の正の整数の組 $(A, B)$ であって、 $A$ の末尾の桁が $B$ の先頭の桁に等しく、 $A$ の先頭の桁が $B$ の末尾に等しいようなものの個数を求めよ。

$1 \leq N \leq 2×10^{5}$

解法

$A = 12345$ とすると、 $B = 5xxxxx1$ となる。(xは存在しないこともある)
じゃあ $A$ を固定して、とりうる間のパターンを数えよう。。。→地獄実装のはじまり\(^0^)/

ここで一度立ち止まって、「先頭が５で末尾が１であるような数はいくつあるか？」 と考える。
これは $N$ を１からなめていけば容易に数えられる。
もっと言えば、 先頭が $i$ で、末尾が $j$ であるような数の個数 は前計算できる！

よって、 $cnt[i\rbrack[j\rbrack = 先頭が i、末尾が j であるような数の個数$ として、 $\sum_{i=0}^{9}\sum_{j=0}^{9} cnt[i\rbrack[j\rbrack × cnt[j\rbrack[i\rbrack$ が答えになる。

実装

N = int(input())
# cnt[i][j] := 先頭がi, 末尾がjで始まるような数
cnt = [[0] * 10 for i in range(10)]
for n in range(1, N + 1):
    cnt[int(str(n)[0])][int(str(n)[-1])] += 1

ans = 0
for i in range(10):
    for j in range(10):
        ans += cnt[i][j] * cnt[j][i]

print(ans)

感想

詰めたと思っても、もっと詰められるかも？という意識を忘れないように。。。。
ABC-Dで筋肉実装なんてそうそう出ないんですよ。。。。反省します