約数の個数がある数より大きい自然数を昇順に列挙する（2）

前回の方法は、約数が小さいほうから順に指数の組合せを生成していきましたが、約数の個数が5万を超えるものを生成するのに2個から始めるというのはかなりムダがあるような気がしますね。実際のところ、ギリギリ5万個を超える指数の組合せ（一つでも指数が減ったら5万以下になる）を列挙するのは簡単です。

def gen_minimum_combinations(n, e0 = 1):
    yield (n,)
    for e in takewhile(lambda e: (e + 1) * e <= n, count(e0)):
        for a in gen_minimum_combinations(n / (e + 1), e):
            yield a + (e,)

5行で書けました。しかし、ここからが難しいのです。このようにして列挙した指数の組合せを値にしたときに昇順に並べなければなりません。最大で2^{50000}という数が出てきてしまいます。こんな計算は避けたいですよね。実際にはPythonはこれくらいの計算は簡単にこなしてしまいますが、できれば多倍長整数は使いたくないところです。これを避ける方法は本当にあれこれ考えたのですが、調べてみると実はここを速くしてもProblem 12を解くのにはたいして速くならないことがわかってしまいました。なので手抜きして対数計算をすることにしました。これなら簡単です。

from itertools import *
from Queue import Queue, PriorityQueue
from math import sqrt, log

def is_prime(n):
    return all(n % p != 0 for p in takewhile(lambda p: p * p <= n, ps))

ps = [ 2 ]
def primes(k):
    if k < len(ps):
        return ps[k]
    
    for n in count(ps[-1] + 1):
        if is_prime(n):
            ps.append(n)
            if len(ps) > k:
                break
    return ps[k]

def value(es):
    return reduce(lambda x, (p, e): x * p ** e,
                            izip(imap(primes, count()), es), 1)

def num_divs(es):
    return reduce(lambda x, e: x * (e + 1), es, 1)

def log_value(es):
    return reduce(lambda x, (p, e): x + log(p) * e,
                            izip(imap(primes, count()), es), 0)

def gen_minimum_combinations(n, e0 = 1):
    yield (n,)
    for e in takewhile(lambda e: (e + 1) * e <= n, count(e0)):
        for a in gen_minimum_combinations(n / (e + 1), e):
            yield a + (e,)

def gen_greater_combinations(N):
    q = Queue()
    for v, es in sorted((log_value(es), es)
                        for es in gen_minimum_combinations(N)):
        q.put(es)
    
    pq = PriorityQueue()
    es_base = q.get()
    pq.put((value(es_base), es_base))
    while not pq.empty():
        v0, es0 = pq.get()
        yield v0, es0
        if es0 == es_base and not q.empty():
            es_base = q.get()
            pq.put((value(es_base), es_base))
        es1 = (es0[0] + 1,) + es0[1:]
        pq.put((value(es1), es1))
        for k in xrange(1, len(es0)):
            if es0[k-1] == es0[k] + 1:
                es2 = es0[:k] + (es0[k] + 1,) + es0[k+1:]
                pq.put((value(es2), es2))
                break
            elif es0[k-1] > es0[k]:
                break
        else:
            if es0[-1] == 1:
                es2 = es0 + (1,)
                pq.put((value(es2), es2))

def is_descending(es):
    return all(e1 >= e2 for e1, e2 in izip(es, es[1:]))

def gen_greater_permutations(N):
    pq = PriorityQueue()
    g = gen_greater_combinations(N)
    set_used = set()
    pq.put(next(g))
    while not pq.empty():
        v, es0 = pq.get()
        yield v, es0
        for k in xrange(len(es0) - 1):
            if es0[k] > es0[k+1]:
                es1 = es0[:k] + (es0[k+1], es0[k]) + es0[k+2:]
                v1 = value(es1)
                if v1 not in set_used:
                    pq.put((v1, es1))
                    set_used.add(v1)
        if is_descending(es0):
            try:
                pq.put(next(g))
            except StopIteration:
                pass

def right_isolated_zero_pos(es):
    for k in xrange(len(es) - 2, 0, -1):
        if es[k] == 0:
            if es[k-1] > 0:
                return k
            else:
                return -1
    return -1

def gen_greater_numbers(N):
    pq = PriorityQueue()
    g = gen_greater_permutations(N)
    pq.put(next(g))
    while not pq.empty():
        v, es0 = pq.get()
        yield v
        
        es1 = es0[:-1] + (0, es0[-1])
        pq.put((value(es1), es1))
        
        pos = right_isolated_zero_pos(es0)
        if pos > 0:
            es2 = es0[:pos-1] + (0, es0[pos-1]) + es0[pos+1:]
            pq.put((value(es2), es2))
        
        if all(e > 0 for e in es0):
            try:
                pq.put(next(g))
            except StopIteration:
                pass

def is_triangle(n):
    m = int(sqrt(n * 2))
    return n == m * (m + 1) / 2

N = 50000
print next(ifilter(is_triangle, gen_greater_numbers(N)))