CodeIQ 「ルート・パワー」問題

CodeIQ の @riverplus さん出題の問題。

の有理数の項を求めよというもの。
他の方のコードが http://togetter.com/li/911507 にある。

とすると、求める数 F(n) は

を解に持つ方程式

を頑張って展開すると

になるから、F(n) は差分方程式

を満たす。
なので、

とすれば

だから

は直接計算して

になる。

下は提出したコードを修正したもの。言語は Ruby。

require 'matrix'

def mod_pow(a, n)
  x = Matrix.I(8)
  loop do
    x = (x * a).map {|y| y % Mod} if n.odd?
    return x if (n >>= 1) == 0
    a = (a ** 2).map {|y| y % Mod}
  end
end

def init_mat
  8.times.map do |i|
    if i < 7 then
      8.times.map {|j| j == i + 1 ? 1 : 0}
    else
      [71, -744, -580, 664, 178, -184, 12, 8]
    end
  end
end

Mod = 10 ** 7
a = Matrix[*init_mat]
v = Vector[1, 1, 11, 31, 285, 1221, 9671, 51171]
p (mod_pow(a, gets.to_i) * v)[0] % Mod

Miller-Rabin 素数判定と Pollard のρ法

最近 Project Euler の問題を解いているのだけど、使い回せるコードを探すのが大変になってきたのでここにまとめておくことにする。
前に他のプログラム中で使ったものだが、とりあえず Miller-Rabin 素数判定と Pollard のρ法(またはモンテカルロ法)による因数分解と、このふたつを使った素因数分解のコードを多少修正したもの。言語は Ruby。

# Miller-Rabin 法による素数判定
class Miller_Rabin
  class << self
    def mod_pow(b, e, n)
      x = 1
      loop do
        x = x * b % n if e.odd?
        return x if (e >>= 1) == 0
        b = b ** 2 % n
      end
    end

    def init_sd
      @s, @d = 0, @n >> 1
      while @d.even? do
        @s += 1
        @d >>= 1
      end
    end

    def witness?(x)
      x = mod_pow(x, @d, @n)
      return true if x == 1 || x == @n - 1
      @s.times.any? do
        x = x ** 2 % @n
        x == @n - 1
      end
    end

    def probably_prime?(n, k = 20)
      @n = n
      return true if @n == 2 || @n == 3
      return false if @n < 2 || @n.even?
      init_sd
      k.times.all? {witness?(rand(@n - 3) + 2)}
    end
  end
end


# Pollard のρ法による因数分解(Brent のアルゴリズムを使用)
class Rho
  class << self
    M = 64

    def f(x)
      (x ** 2 + @c) % @n
    end

    def init_c
      begin
        @c = rand(@n - 1) + 1
      end while @c == @n - 2
    end

    def brent
      k = 1
      x = y = rand(@n)
      loop do
        k.times {x = f(x)}
        m = [k, M].min
        (k / m).times do
          prod = m.times.inject(1) do |q|
            x = f(x)
            q * (x - y) % @n
          end
          g = @n.gcd(prod)
          return g if g > 1
        end
        y = x
        k <<= 1
      end
    end

    # 非自明な因数をひとつ返す
    def factor(n)
      @n = n
      loop do
        init_c
        q = brent
        return q unless q == @n
      end
    end
  end
end


# 素因数分解
class PF
  class << self
    def factorize(n)
      return if n == 1
      return @d[n] += 1 if Miller_Rabin.probably_prime?(n)
      q = Rho.factor(n)
      factorize(q)
      factorize(n / q)
    end

    # Miller-Rabin 法とρ法により素因数分解を求める
    def prime_division(n)
      @d = Hash.new(0)
      factorize(n)
      @d.sort
    end
  end
end


# 以下使用例

# フェルマー数 F_6
F6 = 2 ** (2 ** 6) + 1

# 素数判定
p Miller_Rabin.probably_prime?(F6)

# 証人数を10として素数判定
p Miller_Rabin.probably_prime?(F6, 10)

# 非自明な因数をひとつ見つける
p Rho.factor(F6)

# 素因数分解
p PF.prime_division(F6)

CodeIQ 「ステップ・アップ・サム」問題

CodeIQ の Kawazoe さん出題の問題。
問題と他の方の解答が http://togetter.com/li/875421 にある。

とりあえず項数 1 の数列を許すとして考える。

項数が奇数のとき、m を項数とすると、(m によって定まる)初項 a(m) は

数列が整数列であるための必要十分条件は m が n の約数であること。

項数が偶数のとき、m を 初項+末項 とすると、項数は 、(m によって定まる)初項 b(m) は

数列が整数列であるための必要十分条件は m が n の奇数の約数であること。

初項が正整数となる数列の初項の総和を求めるには、m を n の奇数の約数の範囲で動かし、a(m), b(m) がそれぞれ正のものを選んで足し上げればよい。
m が n の奇数の約数のとき a(m), b(m) は整数、上式より なので、正整数となるのは a(m) と b(m) のうちのどちらか丁度一方で、それは

と書ける。
結局、求める値は

となる。但し m は n の奇数の約数の範囲で動かす。
n を引いているのは項数 1 の数列の寄与を取り除くため。

下が提出したコード。言語は Ruby。
n は標準入力から入力する。
素因数分解を利用しているので、 とかならほとんど計算時間はかからない。

require 'prime'

def f(m, k)
  return (m - 2 * @n / m).abs / 2 + 1 unless @factor[k]
  (0..@factor[k][1]).map {|e| f(m * @factor[k][0] ** e, k + 1)}.inject(:+)
end

@n = gets.to_i
@factor = Prime.prime_division(@n)
p f(1, @n % 2 == 0 ? 1 : 0) - @n

CodeIQ 「スロット・マシン」問題

CodeIQ の Kawazoe さん出題の問題。
問題と他の方のコードが http://togetter.com/li/866620 にある。

をそれぞれ 0〜9 の数字の出現回数とする。
の10個の数を昇順に並べ替えたものを とする。
例えば n=4 のとき の組は次のいずれかになる。

0, 0, 0, 0, 0, 0, 1, 1, 1, 1
0, 0, 0, 0, 0, 0, 0, 1, 1, 2
0, 0, 0, 0, 0, 0, 0, 0, 2, 2
0, 0, 0, 0, 0, 0, 0, 0, 1, 3
0, 0, 0, 0, 0, 0, 0, 0, 0, 4

の中で同じ数が並んでいる部分列をグループと呼ぶことにする。
グループ g で並んでいる数を 、列の長さを と書くことにする。
例えば上の最初の例では と のふたつのグループが存在し、
が与えられたとき、それに対応する組み合わせは

で計算でき( は に含まれるグループについての積)、 に対応する数字の並びが出現する確率はこれを で割ったものになる。
この式は解法pdfの計算法の書き方を変えただけのものなので、証明は省略。

プログラムではメモ化再帰により を、k=10 から始めて k の小さいほうへ向かって生成している( 自体をあらわには扱っていないけど)。再帰1段がグループ1個ぶんの処理。

グループの要素数が取り得る最大値を取る場合、 の組をまとめて計算できる。
例えば n=11 で のとき、値が 4 のグループは要素数が最大値 2。
このとき可能な は

0, 0, 0, 0, 0, 1, 1, 1, 4, 4
0, 0, 0, 0, 0, 0, 1, 2, 4, 4
0, 0, 0, 0, 0, 0, 0, 3, 4, 4

の3通りあるが、これらのいずれかが出現する組み合わせは

とまとめて計算できる。
(例えば 0 と 1 が4個ずつ出現すると決めた時点で、残り3個のリールには 2〜9 の8種の数字がどう出現してもいいとかそういうこと)。こうすることで、全ての の列を列挙するコードと較べて、計算速度が約2倍になる。

下が提出したコードを少し修正したもの。
http://ideone.com/WNELWN に ideone.com での実行結果がある。F(100) だけの計算だと0.16秒だった。

# 着目しているグループの最も大きなインデックスが k
# b_k = b, b1 + b2 + … + b_k = m
def calc(b, m, k)
  @memo[b | m << @shift | k << 2 * @shift] ||= begin
    # r はグループの要素数。まず取り得る最小値とする
    r = [m - k * (b - 1), 1].max
    m -= b * r
    k -= r
    denom = @fac[r] * @fac[b] ** r
    sum = 0
    # このループ1周ごとにグループの要素数 r を1ずつふやしていく
    while m >= b do
      sum += (-(-m / k)...b).inject(0) {|s, c| s + calc(c, m, k)} / denom
      m -= b
      k -= 1
      denom *= (r += 1) * @fac[b]
    end
    # グループの要素数が最大値を取る場合
    # 上で説明したように簡単に計算できる
    sum + @num / (denom * @fac[k] * @fac[m]) * k ** m
  end
end

# F(n) を文字列で返す
def f(n)
  # 階乗  @fac[k] = k!
  @fac = (1..[n, 10].max).each_with_object([1]) {|k, fc| fc << k * fc.last}
  @num = @fac[n] * @fac[10]
  @memo = {}
  @shift = n.to_s(2).size
  (-(-n / 10)..n).inject(0) {|s, b| s + b * calc(b, n, 10)}.
  to_s.insert(-n - 1, '.').sub(/\.?0*$/, '')
end

puts f(6)
puts f(12)
puts f(100)

CodeIQ 「カット・アンド・スクエア」問題

CodeIQ の Kawazoe さん出題の問題。
問題と他の方のコードが CodeIQ「カット・アンド・スクエア」問題 みんなのコード にまとまっている。

とする。
求める n 桁の整数の上下 n/2 桁をそれぞれ x, y とすると

が成り立つ。
変形して


として

つまり を2平方数の和で表す方法を列挙できればよい。
をガウス整数の積に分解して

だから、 のガウス素数への分解が与えられれば、u, v の組は簡単に列挙できる。
このためには、まず を有理素数に素因数分解して、各有理素数をガウス素数に分解すればよい。

例えば n=4 のとき、上の式は

で、 と素因数分解できる。73, 137 をそれぞれ とガウス素数に分解して、

右辺のガウス素数を左辺の因数に適当に割り振る。例えば

とすれば となって、 という解が得られる。
同様に

とすれば となって、 を得る(こちらは自明だけど)。

いちばん下にあるコードが提出したコードを少し修正したもの。言語は Ruby。
n=36 まで計算して ideone.com での実行時間は今走らせたら 0.65秒だった → http://ideone.com/mGfU6N (乱数を使っているので実行時間はばらつきがある)。
提出したものは浮動小数点の平方根 Math.sqrt を使っていたのだが、それだと丸め誤差のせいで大きな n でおかしくなるので、整数の平方根を自前で計算するように修正した。
有理素数への分解は Miller-Rabin 法と Pollard のρ法 (Pollard's rho algorithm)、有理素数のガウス素数への分解は Cornacchia-Smith のアルゴリズムによった。

の有理素数の素因数 q の片方のガウス素数の素因数は q と の共通の因数になる。
Cornacchia-Smith のアルゴリズムのところは(ガウス整数に対して gcd というのはたぶん書き方変だが) をユークリッドの互除方で求めるようなことをしている。
有理素数が4k+3型だと分解できないが、 の素因数は平方剰余の第1補充法則により4k+3型にはならないので、分解できない場合の処理は省略できる。
ユークリッドの互除法を呼んでいるところ

    u = euclid(q, m, sqrt(q))

がちょっと怪しげで、英語版 Wikipedia の記事によるとここは

    r = m % q
    r = q - r if 2 * r > q
    u = euclid(r, q, sqrt(q))

とでもなるところだけど*1、こうでなくてもテストした範囲では大丈夫そうなのでそのままにしてある。
実はここのところを書いたのがかなり前で、そのときは分かっていてこういう書き方になったのかもしれない。

(追記)
やっぱり r=q-r なんて要らないと思う。だって 2r>q なら互除法の最初の1段で r=q-r と同じ計算が行われるではないか。ということで r=m%q も要らなくて r のところにはそのまま m を突っ込んでもいいはず。

# Miller-Rabin法による素数判定
class Miller_Rabin
  def init_sd
    @s, @d = 0, @n >> 1
    while @d.even? do
      @s += 1
      @d >>= 1
    end
  end

  def witness?(x)
    x = mod_pow(x, @d, @n)
    return true if x == 1 || x == @n - 1
    @s.times.any? do
      x = x ** 2 % @n
      x == @n - 1
    end
  end

  def probably_prime?(n, k = 20)
    @n = n
    return true if @n == 2
    return false if @n < 2 || @n.even?
    init_sd
    k.times.all? {witness?(rand(@n - 1) + 1)}
  end
end

# ρ法による因数分解(Brent のアルゴリズムを使用)
class Rho
  M = 64

  def f(x)
    (x ** 2 + @c) % @n
  end

  def init_c
    begin
      @c = rand(@n - 1) + 1
    end while @c == @n - 2
  end

  def factor(n)
    @n = n
    init_c
    k = 1
    x = y = rand(@n)
    loop do
      k.times {x = f(x)}
      m = [k, M].min
      (k / m).times do
        prod = m.times.inject(1) do |q|
          x = f(x)
          q * (x - y) % @n
        end
        g = @n.gcd(prod)
        return g if g > 1
      end
      y = x
      k <<= 1
    end
  end
end

# 素因数分解
class PF
  def factorize(n)
    return if n == 1
    return @d[n] += 1 if Miller_Rabin.new.probably_prime?(n)
    q = Rho.new.factor(n)
    factorize(q)
    factorize(n / q)
  end

  # Miller-Rabin法とρ法により素因数分解を求める
  def prime_division(n)
    @d = Hash.new(0)
    factorize(n)
    @d.sort
  end
end

# b ** e % n を計算
def mod_pow(b, e, n)
  x = 1
  loop do
    x = x * b % n if e.odd?
    return x if (e >>= 1) == 0
    b = b ** 2 % n
  end
end

# ここまで(有理素数への)素因数分解のためのコード
# (ライブラリの素因数分解が速ければ必要ない)

require 'complex'

# 正整数 a に対して floor(sqrt(a)) を求める
def sqrt(a)
  x, y = 1 << a.to_s(4).size, a + 2
  x, y = (x + a / x) / 2, x while x < y
  y
end

# ユークリッドの互除法
def euclid(x, y, c)
  x <= c ? x : euclid(y % x, x, c)
end

# m^2 + 1 をガウス素数の積に分解
def calc_gdiv(m)
  PF.new.prime_division(m ** 2 + 1).flat_map do |q, e|
    # Cornacchia-Smith のアルゴリズムにより有理素数 q をガウス素数に分解
    # m^2+1 が 4k+1 型の素因数しか持たないことと
    # 各素因数 q について m が -1 の mod q での平方根であることを利用して
    # もとのアルゴリズムを簡略化してある
    u = euclid(q, m, sqrt(q))
    v = sqrt(q - u ** 2)
    [u + v * Complex::I] * e
  end
end

# u^2 + v^2 = m^2 + 1 となる非負整数の組 (u, v) を列挙
def div2squares(gdiv, z = gdiv[0], k = 1)
  return [[z.real.abs, z.imag.abs]] if k == gdiv.size
  div2squares(gdiv, z * gdiv[k], k + 1) + div2squares(gdiv, z * gdiv[k].conj, k + 1)
end

# 問題の条件を満たす整数を列挙
def solve(n)
  m = 10 ** (n / 2)
  div2squares(calc_gdiv(m)).flat_map {|u, v|
    u, v = v, u if u.odd?
    [[u, v], [-u, v]]
  }.uniq.map {|u, v| [(u + m) / 2, (v + 1) / 2]}.
  reject {|x, y| x < m / 10 || y <= 1}.
  map {|x, y| m * x + y}.sort
end

# n=36 まで、問題の条件を満たす整数、それらの整数の個数と和を出力
2.step(36, 2) do |n|
  puts "n = #{n}"
  puts a = solve(n)
  puts "n_solutions = #{a.size}\nsum = #{a.inject(0, :+)}\n\n"
end

*1:別のところでの説明ではここの不等号の向きが逆になっていて余計に話がややこしい