「スクエア・カルテット」問題解答 ( CodeIQ )

CodeIQ 数学

はじめに

CodeIQというところでプログラミング問題「スクエア・カルテット」に解答しましたので、ネタバレです。

codeiq.jp

問題

標準入力から、空白区切り1行で与えられる自然数$a,b ( 1\leq a\le b\le 10^5 )$に対し、以下で定義される$F(a,b)$を計算して出力するという問題です。

  • $F(a,b)$:$x^2+a^2=y^2+b^2$の自然数解$(x,y)$全ての組み合わせに対する$x+y$の合計 ( x,yが逆転しているだけの組は除く )

数式としての表現は、

  • $F(a,b)=\sum_{x^2+a^2=y^2+b^2,x\geq y}(x+y)$ ( x,yは自然数 )

問題で例として挙げられていたのは、$(a,b)=(3,10)$に対して、$x^2+a^2=y^2+b^2$の自然数解は$(x,y)=(10,3),(46,45)$ ( x,yを逆順にした解は除くので ) ということから、$F(3,10)=10+3+46+45=104$ です。

解説

アプローチ

まず、x,yの方程式を変形して、

  • $(x-y)(x+y)=b^2-a^2$

というところから、$(x-y),(x+y)$はそれぞれ$b^2-a^2$の約数です。
加えて、$(x-y),(x+y)$の偶奇が一致することにも着目すると、

  • $pq=b^2-a^2$ ( ただし $p\equiv q\,mod 2,\,0\lt p\lt q$ ) に対し、解$(x,y)=(\frac{p+q}{2},\frac{q-p}{2})$が1対1対応する

ということです。このとき、$x+y=q$ですから、結局、

  • $F(a,b)=\sum_{pq=b^2-a^2, p\equiv q\,mod 2,\,0\lt p\lt q} q$
    $\Leftrightarrow F(a,b)=\sum_{b^2-a^2\equiv 0 \,mod \,p,\,0\lt p\lt \sqrt{b^2-a^2},\,p\equiv \frac{b^2-a^2}{p} \,mod \,2} \frac{b^2-a^2}{p}$

と、約数の組 ( ただし偶奇が一致するもの ) をあげて、大きい方の合計を計算すればよいと分かります。
問題の例であった$F(3,10)$であれば、$10^2-3^2=91=1\times 91=7\times 13$であることから、$F(3,10)=91+13=104$と計算できます。

golf版実装

提出版は以下のRuby(87)でした。( 1行にまとめていましたが同じことです )

a,b=gets.split.map &:to_i
n=b*b-a*a
x=i=0
n%i<1&&n/i-i&1<1&&x+=n/i while(i+=1)**2<n
p x

まあでもせめて、次のRuby(82)にはできましたね。

n=-eval(gets.gsub(/\s/,"**2-")+?0)
x=i=0
x+=n/i*1[n%i|n/i-i&1]while(i+=1)**2<n
p x

別の実装

実は、最初はRubyではなくPerlで組んでいたのですが…。

まず、ループを逆 ( 大きい数から減らしていく方向 ) に書いてTLE、次に通常の整数に収まらないケース ( $F(19126,98765)=16043311668$ ) のケアがなくてWA、bigintを使って任意精度整数対応したら遅くてTLEというマヌケな結果に…。

尤も、通常の整数に収まらなくて困るのはRuby版と同じようにビット演算で &1 を使ったからで、次のPerl(73)のように %2 と剰余演算にしておけばよかったんですけど。内部的には浮動小数点で扱ってくれるので、今回の問題であれば一応精度が足りるのですね。

$/=$";$n=-<>**2+<>**2;$n%$i|($n/$i-$i)%2or$_+=$n/$i while++$i**2<$n;print

ともあれ、Perlで「入力は非bigint」「答えはbigint」として高速化したのが次のコードです。

これは、$b^2-a^2=(b+a)(b-a)$と分解できることから、$b+a,b-a$それぞれを素因数分解して約数を調べるようにしたものです。…まあ、素因数分解で特に 工夫はしてないですが。

終わりに

実はこれ、今年一発目に解いた問題でした。今年も面白い問題をよろしくお願いします、ということで。