Project Euler 119

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


調べる範囲を狭めることができます。例えば2乗のとき、3桁の2乗は6桁以下だから各桁を足しても54以下です。2桁なら2乗して4桁以下で各桁を足して36以下だから、底が36より大きいことはありません。
各べき乗についてジェネレータを作り、範囲を区切って出したものをまとめて昇順にソートして出します。
最初に定義しているクラスは無駄な計算を避けるためのものです。区切るためにtakewhileを使うと、最初に叙述関数がFalseになったときにその項が失われてしまいます。なので、最初にFalseになったところまで出すtakewhile2を作ったのですが、それだとたいていの場合いきなりFalseになってしまい、どんどん大きな数が計算されてしまいます。それを避けるために一度計算してもその値を残しておくようなtakewhileを定義したクラスを作りました。

from itertools import *

class cGen:
    def __init__(self, g):
        self.g = g
        self.n = -1
    
    def takewhile(self, pred):
        if self.n != -1:
            if pred(self.n):
                yield self.n
                self.n = -1
            else:
                raise StopIteration
        
        for self.n in self.g:
            if pred(self.n):
                yield self.n
                self.n = -1
            else:
                raise StopIteration
        raise StopIteration

def gen_powers(e):
    def max_base(e):
        return next(9 * (m - 1) * e for m in count(2)
                            if 9 * m * e < 10 ** (m - 1))
    
    return ((n, n ** e) for n in xrange(2, max_base(e) + 1))

def gen_pows():
    def sum_digits(n):
        return 0 if n == 0 else n % 10 + sum_digits(n / 10)
    
    a = []
    gs = []
    for e in count(2):
        limit = 2 ** (e + 1)
        gs.append(cGen(m for n, m in gen_powers(e)
                                if sum_digits(m) == n))
        for g in gs:
            a.extend(g.takewhile(lambda m: m <= limit))
        a.sort()
        if len(a) > 0:
            for k in takewhile(lambda k: a[0] < limit, xrange(len(a))):
                yield a.pop(0)

N = 30
print next(n for k, n in izip(count(1), gen_pows()) if k == N)