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

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

$\Large m = 10^{\frac{n}{2}}$
とする。
求める n 桁の整数の上下 n/2 桁をそれぞれ x, y とすると
$\Large x^2 + y^2 = mx + y$
が成り立つ。
変形して
$\Large (2x-m)^2 + (2y-1)^2 = m^2 + 1$
$\Large u=2x-m,\; v=2y-1$
として
$\Large u^2 + v^2 = m^2 + 1$

つまり $m^2+1$ を2平方数の和で表す方法を列挙できればよい。
$u^2 + v^2$ をガウス整数の積に分解して
$\Large (u + {\rm i}v) (u - {\rm i}v) = m^2 + 1$
だから、 $m^2+1$ のガウス素数への分解が与えられれば、u, v の組は簡単に列挙できる。
このためには、まず $m^2+1$ を有理素数に素因数分解して、各有理素数をガウス素数に分解すればよい。

例えば n=4 のとき、上の式は
$\Large (u+{\rm i}v)(u-{\rm i}v) = 10001$
で、 $10001 = 73\cdot 137$ と素因数分解できる。73, 137 をそれぞれ $73 = (8+3{\rm i})(8-3{\rm i}),\; 137 = (11+ 4{\rm i})(11-4{\rm i})$ とガウス素数に分解して、
$\Large (u+{\rm i}v)(u-{\rm i}v) = (8+3{\rm i})(8-3{\rm i})(11+ 4{\rm i})(11-4{\rm i})$
右辺のガウス素数を左辺の因数に適当に割り振る。例えば
$\Large u+{\rm i}v = (8+3{\rm i})(11+ 4{\rm i}),\; u-{\rm i}v = (8-3{\rm i})(11-4{\rm i})$
とすれば $u = 76,\; v = 65$ となって、 $76^2+65^2 = 10001$ という解が得られる。
同様に
$\Large u+{\rm i}v = (8+3{\rm i})(11- 4{\rm i}),\; u-{\rm i}v = (8-3{\rm i})(11+4{\rm i})$
とすれば $u = 100,\; v = 1$ となって、 $100^2+1^2 = 10001$ を得る(こちらは自明だけど)。

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

$m^2+1$ の有理素数の素因数 q の片方のガウス素数の素因数は q と $m+{\rm i}$ の共通の因数になる。
Cornacchia-Smith のアルゴリズムのところは(ガウス整数に対して gcd というのはたぶん書き方変だが) ${\rm gcd}(q, m+{\rm i})$ をユークリッドの互除方で求めるようなことをしている。
有理素数が4k+3型だと分解できないが、 $m^2+1$ の素因数は平方剰余の第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:別のところでの説明ではここの不等号の向きが逆になっていて余計に話がややこしい