Rubyで高速素因数分解したい (ABC177E)

参考

https://atcoder.jp/contests/abc177/editorial/82
(これを読んでからじゃないと記事が意味不明だと思います)

本題

A以下の数がたくさん与えられて、それらを素因数分解したいとき、
エラトステネスの篩を行って、ついでにふるい落とした数を入れておくと、
素因数分解のときに楽

cache = {key(その値) => value(keyをふるい落とした最小の素因数) }

みたいなcacheを作ると楽だよという話です。

例) 100を素因数分解する場合

cache[100] 
#=> 2 = 2で割れる
cache[100/2]
#=> 2 = 2で割れる
cache[50/2] 
#=> 5 = 5で割れる
cache[25/5] 
#=> 5 = 5で割れる

#∴ 100 は 2で2回、5で2回割れる

このcacheを作っておけば、素因数分解をO(logN)でできるねという話です。
今回はこれを用いて素因数分解を高速にできるクラスを作ってみたので、
既存のInteger#prime_division との性能比較をしたいと思います。

作ったもの

ソースコード

# エラトステネスの篩で素数列挙して高速に素因数分解したい！
# ref: https://atcoder.jp/contests/abc177/editorial/82
class PrimeDivWithSieve
  def initialize(n)
    @sieve = [] # nまでの素数を入れる
    @min_div = {} # keyの値の最小の素因数を入れる
    # 他を篩落とし得る素数はsqrtを上限にできる
    (2..Math.sqrt(n).floor).each do |i|
      next if @min_div[i] # ここに値が入ってる = ふるい落とされている
      @sieve << i # ふるい落とされずに来たらそいつは素数

      sieve_target = i * i
      while sieve_target <= n do
        @min_div[sieve_target] ||= i
        sieve_target += i
      end
    end
    (Math.sqrt(n).floor.next..n).each do |i|
      next if @min_div[i]
      @sieve << i
    end
  end

  # Integer#prime_division と同じ値を返すようにする
  # https://docs.ruby-lang.org/ja/latest/method/Integer/i/prime_division.html
  def prime_division(num)
    return [[num, 1]] if !@min_div[num] # 素数のときすぐ返す

    return_array = [] # [[a, x], [b, y]] <=> num = a^x * b^y
    while num > 1 do
      prime = @min_div[num] # 最小の素因数, nil => numが素数
      break return_array.push([num, 1]) if !prime

      div_total = 0
      while num % prime == 0 do
        num /= prime
        div_total += 1
      end
      return_array.push([prime, div_total])
    end

    return_array
  end

  def prime_list
    @sieve
  end
end

性能比較

(テストコード全文)[https://github.com/k-karen/ruby-samples/blob/master/sieve/bench_prime_division.rb]

N = 1_000_000
times = N/25 # テスト件数
TESTCASES = (1..N).to_a.sample(times)
require 'prime'

ans1 = []
ans2 = []
Benchmark.bm 10 do |r|
  r.report 'MyDivsor' do
    divisor = PrimeDivWithSieve.new(N)
    TESTCASES.each do |i|
      ans1.push(divisor.prime_division(i))
    end
  end

  r.report 'PrimeDiv' do
    TESTCASES.each do |i|
      ans2.push(i.prime_division)
    end
  end
end

# times = N/25 (40,000件)
# Result
# user     system      total        real
# MyDivsor     0.875262   0.032392   0.907654 (  0.926605)
# PrimeDiv     0.849263   0.012468   0.861731 (  0.879886)

# times = N/2 (500,000件)
# Result
# user     system      total        real
# MyDivsor     1.659268   0.058786   1.718054 (  1.758668)
# PrimeDiv    10.787444   0.118755  10.906199 ( 11.071594)

件数少ないと微妙ですが、件数増えると篩利用してるほうが早そうです。

ABC177E

https://atcoder.jp/contests/abc177/submissions/16390851
これで勝つると思ってsubmitしたんですが、1つTLEがどうしても消えませんでした。
今回は何回割れるかという情報がいらないので、
もう篩作るときにその値が持ってる素因数をメモしておけばいいじゃんということになりました。
例) 100 => [2,5], 99 => [3, 11]

https://atcoder.jp/contests/abc177/submissions/16391235
これで通りました。
(今回の高速素因数分解とは違うろじっくなのでおまけ程度にどうぞ)

nまでの各値の素因数を計算(with篩)
https://github.com/k-karen/ruby-samples/blob/master/sieve/enum_elments.rb

ソースコード

# 素因数だけをnまで計算する
class EnumElements
  def initialize(n)
    @sieve = []
    @elements = {}
    (2..n).each do |i|
      next if @elements[i]
      @sieve << i
      @elements[i] = [i]

      sieve_target = i * 2
      while sieve_target <= n do
        if @elements[sieve_target]
          @elements[sieve_target].push(i)
        else
          @elements[sieve_target] = [i]
        end
        sieve_target += i
      end
    end
  end

  def elements(num)
    @elements[num] || []
  end
end