Ruby と crystal で解く AtCoder ABC254 D

はじめに

AtCoder さん、ありがとうございます。

D - Together Square

時々出題される、素数系の問題。

愚直な例（失敗）

n = gets.to_i
ans = 0
a = []
k = 1
while k <= n
  a << k * k
  k += 1
end
k = 1
while k <= n
  a.each do |x|
    ans += 1 if x >= k && x % k == 0 && x / k <= n
  end
  k += 1
end
puts ans

while k <= n
  a.each do |x|

ここの二重ループが $O(N^2)$ となってしまいますので、TLEとなります。

Integer#prime_division

Ruby

require "prime"

def calc_odd_factor_prod(n)
  ret = 1
  n.prime_division.each do |k, v|
    ret *= k if v.odd?
  end
  ret
end

n = gets.to_i
ans = 0
(1..n).each do |i|
  odd = calc_odd_factor_prod(i)
  j = 1
  while odd * j * j <= n
    ans += 1
    j += 1
  end
end
puts ans

素数とくればrequire "prime"
Integer#prime_divisionで素因数とその指数を算出します。

crystal

crystalには標準で素数関連のクラスが準備されていませんので、自作もしくは公開されているコードを使用する必要があります。

素数関連クラスがあるRubyの方が珍しいかも知れませんね。

	Ruby	crystal
実行時間(ms)	1630	not started

u2dayo さんコード

いつも参考にさせていただいております、@u2dayo さんのコードを写経移植してみました。

Ruby

def calc_odd_factor_prod(n)
  i = 2
  while i * i <= n
    i2 = i * i
    while n % i2 == 0
      n /= i2
    end
    i += 1
  end
  n
end

n = gets.to_i
ans = 0
(1..n).each do |i|
  odd = calc_odd_factor_prod(i)
  (1..n).each do |j|
    break if odd * j * j > n
    ans += 1
  end
end
puts ans

残念ですがTLE、後200ms位改善すればACすると思われます。
本問題は、RubyユーザでのACが少なかったことも頷けます。

crystal

def calc_odd_factor_prod(n)
  i = 2
  while i * i <= n
    i2 = i * i
    while n % i2 == 0      
      n //= i2
    end
    i += 1
  end
  n
end

n = read_line.to_i
ans = 0
(1..n).each do |i|
  odd = calc_odd_factor_prod(i)
  (1..n).each do |j|
    break if odd * j * j > n
    ans += 1
  end
end
puts ans

計算量が多くても大丈夫。そう、crystalizer2ならね。

	Ruby	crystal
実行時間(ms)	TLE	199

radiol さんコード

@radiol さんのコードも写経移植してみました。

Ruby

n = gets.to_i
is_sqrt = Array.new(n + 1, false)
(1..).each do |i|
  break if i * i > n
  is_sqrt[i * i] = true
end
factor = Array.new(n + 1){ [] }
(1..n).each do |i|
  i.step(n, i) do |j|
    factor[j] << i
  end
end
cnt = Hash.new(0)
(1..n).each do |i|
  f = 0
  factor[i].reverse.each do |x|
    if is_sqrt[x]
      f = x
      break
    end
  end
  cnt[i / f] += 1
end
ans = 0
cnt.values.each do |v|
  ans += v * v
end
puts ans

factor配列を用意し、大きい順で検索をしている工夫が光ります。

crystal

n = read_line.to_i
is_sqrt = Array.new(n + 1, false)       
(1..).each do |i|
  break if i * i > n
  is_sqrt[i * i] = true
end
factor = Array.new(n + 1){ [] of Int32 }
(1..n).each do |i|
  i.step(to: n, by: i) do |j|
    factor[j] << i
  end
end
cnt = Hash(Int32, Int32).new(0)
(1..n).each do |i|
  f = 0
  factor[i].reverse.each do |x|
    if is_sqrt[x]
      f = x
      break
    end
  end
  cnt[i // f] += 1
end
ans = 0
cnt.values.each do |v|
  ans += v * v
end
puts ans

crystalのstepメソッド、初めて使用しました。

	Ruby	crystal
実行時間(ms)	643	201

気づき（ceil/floor）

Rubyのceilやfloorは整数を返しますが、crystalの場合浮動小数点数を返します。