【Ruby】.countメソッドの裏側 ~競プロでも使える桁和計算で応用~

例題

正の整数nの桁和を、「nを十進法で表したときの各桁の和」と定める。
例えば、2026の桁和は 2 + 0 + 2 + 6 = 10である。

入力が
n k
で与えられるとき、
N以下の正整数のうち桁和がKであるものの個数を求めよ。

入力例:

30 4

出力例:

3

30以下の整数で桁和が4になるものは、4, 13, 22の3つ。
　

引用元: https://atcoder.jp/contests/abc444/tasks/abc444_b

解法

この問題の解き方としては

1. ゴリ押し: 1からnまでの整数でループを回し、各桁の総和がKと一致したらcount+=1
2. 動的計画法: dp[i][j]に1~iまでで桁和がjになるものの数を順次入れ、答えはdp[n][k]

の2つが考えられるのですが、ここでは1.ゴリ押しの解法をよりスマートに書くrubyのメソッドチェーンの考え方を活用してみましょう。

ゴリ押し

フツーにゴリ押しするならこんな感じ:

# 入力を受け取る
n, k = gets.split(' ').map(&:to_i)

# 1からnまででループを作り、桁の合計を算出する
count = 0
(1..n).each do |num|
  sum = 0
  while num > 0
    # numを10で割った余り（＝一桁目の数字）をsumに加える
    # 例: num=259の場合、num%10=9より一桁目の9が取り出せる
    sum += (num % 10)

    # 桁を落とすためにnumを10で割った商をnumに再代入する
    # 例: num=259の場合、num/10＝25、num/10=2、num/10=0とどんどん桁が落ちていく
    num /= 10
  end

  # もし桁和sumがkと一致する場合はcount+=1
  if sum == k
    count += 1
  end
end

puts count

愚直にやると、こうして10で割った余りと商をうまく駆使することで回答できます。

もっとスマートなゴリ押し

ただこのロジックだと、一見何をしているのか分かりづらく、いかにも実務などであまり好まれなさそうな感じ。

そこでもっとスマートに書けるのがこれです。

# 入力を受け取る
n, k = gets.split(' ').map(&:to_i)

(1..n).count do |num|
  # num=259.to_s → "259"
  # "259".chars → ["2", "5", "9"]
  # ["2", "5", "9"].map(&:to_i) → [2, 5, 9]
  # [2, 5, 9].sum → 16
  num.to_s.chars.map(&:to_i).sum == k
end

puts count

これであれば、長くて分かりづらかったコードもかなり簡潔にまとめられます。
（計算効率は多少落ちますが、、、）

countメソッドの裏側

このコードを理解する上では、countメソッドの仕組みを理解する必要があります。

(1..n).count do |num|
  num.to_s.chars.map(&:to_i).sum == k
end

この部分ですが、doブロックの中が条件式になっており、true/falseを返すようになっています。
countメソッドは、内部的にcount=0を自分で初期値設定し、ループを回す過程で条件式がtrueになればcountを1増やすという仕組みで動いています。
　
これにより、内部の条件式
「桁和がkと等しい」
がtrueの時だけcountが増えていくようになっています。
　

実は普段使うようなcount、例えば下のコード

["apple", "orange", "banana", "orange", "lemon"].count("orange")
# => 2

でも、実は内部的な処理をちゃんと書き起こすと

["apple", "orange", "banana", "orange", "lemon"].count do |fruit|
  fruit == "orange"
end
# => 2

のように中で比較条件式が動いており、trueであればcountがインクリメントされるようになっています。

まとめ

こういう感じで、メソッドの裏側の動きを理解できると面白いし、書ける人かっこいいですよね〜

AtCoderをはじめ、競プロでは一つの問題に対して色んな解法があるため、
それらを比較して、「なぜこの解き方が一番いいか」「もっとこうしたら簡潔に書けるのでは」「多少長いけどこっちの方が読みやすいのでは」などとあれこれ考える時間が幸せです。

結局何を伝えたいんだ、という感じになりましたが
ぜひコードの理解を深めるきっかけとして競プロを活用してみてはいかがでしょうか！

「意外と知らなかった...」
「競プロやってみようかな...」
と少しでも思ってくれた方は、どうか！清き1いいねをお願いします！🥹

読んでくださりありがとうございました！