【Project Euler】Problem 64: 平方根の連分数

• 本記事はProjectEulerの「100番以下の問題の説明は記載可能」という規定に基づいて回答のヒントが書かれていますので、自分である程度考えてみてから読まれることをお勧めします。

問題：平方根の連分数の循環周期

原文 Problem 64:Odd period square roots

問題の要約：　$1<N<10^4$で$\sqrt{}N$の連分数の循環周期が奇数になるもの数を求めよ

連分数に関してはこちらに説明があるので参考にしてください。
正の整数の平方根の連分数展開は循環することが知られていますが、sympyに2次無理数の連分数展開を求めるcontinued_fractionという関数があり（説明（英文）はこちら）以下のように簡単に求めることが出来ます。
従ってこの循環部分(=cf[1])の長さをチェックすればよいのですが残念ながらこの関数は遅すぎて$1<N<10^3$で2分程度かかってしまいました。

from sympy.ntheory.continued_fraction import continued_fraction
from sympy import sqrt

for n in range(2,10+1):
  cf = continued_fraction(sqrt(n))
  print(n, cf)
#2 [1, [2]]
#3 [1, [1, 2]]
#4 [2]
#5 [2, [4]]
#6 [2, [2, 4]]
#7 [2, [1, 1, 1, 4]]
#8 [2, [1, 4]]
#9 [3]
#10 [3, [6]]

そこで色々探した結果こちらの「平方根の連分数とペル方程式 有澤健治著」という文献に添付されているコードを参考させていただき以下の関数cfsqrtを書きました。機能はcontinued_fractionと同等ですが平方根の連分数に特化しているのでかなり速く$1<N<10^4$でも1秒以内に答えをだすことが出来ました。

def cfsqrt(m):
  n0 = int(m**(1/2))
  if n0*n0 == m: return [n0]
  n,a,b,cf = n0,1,0,[]
  while True:
    b = n*a -b
    a = (m-b*b)//a
    n = (n0+b)//a
    cf.append(n)
    if a == 1: break
  return [n0,cf]