Project Euler: Problem 44

問題はこちら。5角数(と日本語で言うのか知らないけど)Pi, Pj(i

def main_guess():
    penta_iter = (n*(3*n-1)//2 for n in itertools.count(1))
    penta_set = set()    
    
    for x in penta_iter:
    'x:=Pi+Pjとし, xを順に大きくして探索する.'
        penta_set.add(x)
        for y in (y for y in penta_set if y < D):
            'y:=Pj-Piとし, yも順に大きくする.'
            if (x+y)/2 in penta_set and (x-y)/2 in penta_set:
                return y

確かにこれは正解を出すのだが、はじめに見つかったyが本当に最小である保証がない(実は最小なのだけれど)。もっと後にあるxに対して、より小さなyがある可能性がある。その可能性がないということを証明するのは、そんなに自明ではないはずだ。

厳密に解を求めようとして、自分が書いたのが次のコード。

def is_square(x):
    return math.floor(math.sqrt(x))**2 == x

def is_pentagonal(x):
    return is_square(24*x + 1) and math.floor(math.sqrt(24*x + 1)) % 6 == 5

def main():
    penta_iter = (n*(3*n-1)//2 for n in itertools.count(1))
    penta_set = set()
    D = float('inf')
    
    for j, Pj in enumerate(penta_iter):
        if 3*j - 2 >= D: break #Pj - Pj-1 >= D よりこれ以上の探索は不要.
        penta_set.add(Pj)
        for Pi in (Pi for Pi in penta_set if Pj - Pi < D):
            if is_pentagonal(Pi + Pj) and Pj-Pi in penta_set:
                D = min((D, Pj-Pi))
    return D

Dを更新して探索範囲を絞っていくので、必ず止まることは分かっている。しかしながら、最適解のDがかなり大きい数なので、時間がかかりすぎる。