Solving Google Treasure Hunt Puzzle 4: Prime Numbers

Looks like a nice showcase for Python's stream-processing approach via generators and other iterators/itertools/etc. Here's the relevant snippet from the itoo.py in my personal toolbox:

def sieve():
    """Sieve of Eratosthenes, yield each prime in sequence."""
    yield 2
    D = {}
    q = 3
    while True:
        p = D.pop(q, 0)
        if p:
            x = q + p
            while x in D: x += p
            D[x] = p
        else:
            yield q
            D[q*q] = 2*q
        q += 2

def running_sum(it, k):
    """Yield sums of the latest k items from it"""
    n = it.next
    d = collections.deque(n() for i in range(k))
    s = sum(d)
    while True:
        yield s
        s -= d.popleft()
        d.append(n())
        s += d[-1]

def inters(*its, **k):
    """Yield items that are in ALL of sorted iterators *its"""
    heap = []
    for it in its:
        try: heap.append((it.next(), it.next))
        except StopIteration: return
    heapq.heapify(heap)
    mx = max(h[0] for h in heap)
    while True:
        v, n = heapq.heappop(heap)
        if v==mx: yield v
        try: v = n()
        except StopIteration: return
        if v>mx: mx = v
        heapq.heappush(heap, (v, n))

[of course there are several other parts of itoo.py but I've picked only the ones used in the following]. Now, using these tools I keep lying around, I wrote uit.py:

import itertools
import itoo

def prunsums(*totals):
    ps = itertools.tee(itoo.sieve(), len(totals)+1)
    sums = [itoo.running_sum(p, t) for p, t in zip(ps, totals)]
    return itoo.inters(ps[-1], *sums)

def main():
    gen = prunsums(3, 6)
    print gen.next()
    gen = prunsums(7, 17, 41, 541)
    print gen.next()

main()

and:

alex-martellis-macbook-air:~ alex$ time python uit.py
41
7830239

real	0m15.552s
user	0m14.541s
sys	0m0.440s

Of course, it IS only a Macbook Air, so a faster machine would get better times than 15 seconds!-) But, 15 seconds IS about the additional time it took to write uit.py, and I'd say that's about par for the course.

Alex