Project Euler 140

http://projecteuler.net/index.php?section=problems&id=140

Problem 137と同じように見えますが、より難しくなっています。
同じように計算していくと、

(n + 3)x² + (n + 1)x - n = 0

xが有理数になる必要十分条件は判別式が平方数になることだから、

5n² + 14n + 1 = m²
(5n + 7)² - 5m² = 44

ペル方程式に似た方程式、

x² - 5y² = 44

が解ければよいです。一般的に解法がどうなっているのか知りませんが、(x₁, y₁)をx² - 5y² = ±11の解、(x₂, y₂)をx² - 5y² = ±4の解（符号同順）とすれば、x + √5y = (x₁ + √5y₁)(x₂ + √5y₂)で与えられる(x, y)が元の解になっていることが簡単に確かめられます。後者はペル方程式なので、前者を求めます。

x² - 5y² = 11

最小解は(4, 1)なので、これにx² - 5y² = 1の最小解(9, 4)を「掛けて」いけばよい、と思いきやこれでは解が足りないんですね。別にx² - 5y² = -11の最小解(3, 2)の系統の解があります。これに、x² - 5y² = -1の最小解(2, 1)を「掛けた」(16, 7)から始まる解の列があります。同様に負のほうも求めて、±4の解と「掛けます」。昇順に出すのが難しいので、それは範囲を区切って適当に。あと重複があるので、それを排除します。

from itertools import *
from math import sqrt

def unique(g):
    prev = next(g)
    for e in g:
        if e != prev:
            yield prev
        prev = e
    yield prev

def unfold(f, s):
    while True:
        x, s = f(s)
        yield x

def gen_Pells(a, r, n):
    def mul(p, q):
        x1, y1 = p
        x2, y2 = q
        return (x1 * x2 + n * y1 * y2, x1 * y2 + y1 * x2)
    
    return unfold(lambda s: (s, mul(s, r)), a)

def gen_Pells4(a, r, n):
    def mul(p, q):
        x1, y1 = p
        x2, y2 = q
        return ((x1 * x2 + n * y1 * y2) / 2, (x1 * y2 + y1 * x2) / 2)
    
    return unfold(lambda s: (s, mul(s, r)), a)

def combine(g1, g2):
    v1 = g1.next()
    v2 = g2.next()
    while True:
        if v1 == v2:
            yield v1
            v1 = g1.next()
            v2 = g2.next()
        elif v1 < v2:
            yield v1
            v1 = g1.next()
        else:
            yield v2
            v2 = g2.next()

def gen_Pells_11():
    return combine(gen_Pells(( 4, 1), (9, 4), 5),
                   gen_Pells((16, 7), (9, 4), 5))


def gen_Pells__11():
    return combine(gen_Pells(( 3, 2), (9, 4), 5),
                   gen_Pells((13, 6), (9, 4), 5))

def gen_Pells_44():
    def mul(p, q):
        x1, y1 = p
        x2, y2 = q
        return (x1 * x2 + 5 * y1 * y2, x1 * y2 + y1 * x2)
    
    for y_limit in unfold(lambda s: (s, s * 10), 1):
        a = [ ]
        for p in takewhile(lambda p: p[1] < y_limit * 10, gen_Pells_4()):
            for r in takewhile(lambda r: r[1] <= y_limit * 10,
                     dropwhile(lambda r: r[1] <  y_limit,
                            (mul(p, q) for q in gen_Pells_11()))):
                a.append(r)
        
        for p in takewhile(lambda p: p[1] < y_limit * 10, gen_Pells__4()):
            for r in takewhile(lambda r: r[1] <= y_limit * 10,
                     dropwhile(lambda r: r[1] <  y_limit,
                            (mul(p, q) for q in gen_Pells__11()))):
                a.append(r)
        
        for r in unique(r for r in sorted(a)):
            yield r

N = 30
def gen_Pells__4(): return gen_Pells4((1, 1), (3, 1), 5)
def gen_Pells_4():  return gen_Pells4((3, 1), (3, 1), 5)
g = ((n - 7) / 5 for n, m in gen_Pells_44() if n % 5 == 2)
print(sum(islice(g, N)))