AtCoder ABC 152 D - Handstand 2 (400 点)

2回目のABC

問題へのリンク

問題概要

正の整数Nが与えられる。N以下の整数の組（A, B）について、

• Aの末尾の桁とBの先頭の桁が等しい、かつ、Aの先頭の桁とBの末尾の桁が等しい

となるものの個数を求める。ただしA, Bは先頭に0をつけない10進数表記である。

制約

• 1 <= N <= 2x10⁵

解法1（本番でやったTLE）

やり方が思いつかなくてとりあえずサブミットしてみたやつ。 O(N²)なのでサンプルテストですでにTLEしてた笑

n = gets.to_i
count = 0

(1..n).each do |i|
  (1..n).each do |j|
    count += 1 if i.to_s[0] == j.to_s[-1] && i.to_s[-1] == j.to_s[0]
  end
end

puts count

解法2（解説を参考に）

本番ではペアを作ることは何となく思いついたが、それをセットにして数え上げておくところまで至らなかった。桁の違いを考慮しなければと思っていたが、そういう訳でもなかった。

まず、Nを先頭の桁と末尾の桁（Nが1桁ならそれぞれ1桁目）の組に変換したペアが、それぞれ何個ずつあるか数えおく。

1からNまでの整数について、ペアの対となるペアを数え上げる。O(N)になる。

n = gets.to_i

def encoder(x)
  a = x.to_s[0]
  b = x.to_s[-1]
  
  a+b
end

map = {}
(1..n).each do |i|
  encoded = encoder(i)
  if map[encoded]
    map[encoded] += 1
  else
    map[encoded] = 1
  end
end

count = 0
(1..n).each do |j|
  encoded = encoder(j)
  oposite = encoded[1]+encoded[0]
  
  count += map[oposite] if map[oposite]
end

puts count

確認したところ構造体をハッシュのキーにできるので、encoderの返り値は(a, b)的な構造体にしてもよい。