Project Euler 196

プロジェクトオイラー
http://projecteuler.net/index.php

Q196.
自然数を三角形状に並べる。3つの素数の組がそのうち1つを選ぶと残りの2つが隣（斜め含む）になっているときprime tripletと呼ぶことにする。S(n)をn列上の素数でprime tripletの要素の一つになっているものの総和とする。S(5678027) + S(7208785)を求めよ。

あれこれ考えたのだが、結局その列を含む5列をエラトステネスのふるいにかけて素数をあぶりだすのが速かった。

from math import sqrt

import time

from itertools import count
def is_prime(n):

    for p in primes:

        if p * p > n:

            return True

        elif n % p == 0:

            return False

    return True
def make_primes(n):

    m = int(sqrt(n + 0.5))

    for p in xrange(3, m + 1, 2):

        if is_prime(p):

            primes.append(p)

    

    a = [ 0 ] * ((n + 15) / 16)

    for p in primes:

        if p * p > n:

            break

        if p == 2:

            continue

        for k in xrange(3 * p, n, 2 * p):

            a[k>>4] |= 1 << (k & 15)

    

    for k in xrange((m + 1) / 2 * 2 + 1, n + 1, 2):

        if (a[k>>4] & (1 << (k & 15))) == 0:

            primes.append(k)
def add(a, b, x, y):

    if a[y][x]:

        b.append( (x, y) )
def sieve(n):

    a = [ [ True ] * k for k in range(n - 2, n + 3) ]

    for p in primes:

        for k in range(n - 2, n + 3):

            mod = ((k - 1) * k / 2 + 1) % p

            if mod == 0:

                start = 0

            else:

                start = p - mod

            for l in xrange(start, k, p):

                a[k-n+2][l] = False

    return a
def triplet(a, x, y, n):

    b = [ ]

    if x > 0:

        add(a, b, x - 1, y - 1)

        add(a, b, x - 1, y)

        add(a, b, x - 1, y + 1)

    add(a, b, x, y - 1)

    add(a, b, x, y)

    add(a, b, x, y + 1)

    if x < n + y - 4:

        add(a, b, x + 1, y - 1)

    add(a, b, x + 1, y)

    add(a, b, x + 1, y + 1)

    

    return b
def pt_to_num(pt, n):

    k = pt[1] + n - 2

    return (k - 1) * k / 2 + 1 + pt[0]
def S(n):

    m = [ k * (k + 1) / 2 + 1 for k in range(n - 3, n + 3) ]

    a = sieve(n)

    

    s = set()

    for y in range(1, 4):

        for x in range(len(a[y])):

            if a[y][x]:

                b = triplet(a, x, y, n)

                if len(b) >= 3:

                    for pt in filter(lambda pt: pt[1] == 2, b):

                        s.add(pt_to_num(pt, n))

    return sum(s)
t = time.clock()

N = 5678027

M = 7208785

L = int( (max(N, M) + 3) / sqrt(2))

primes = [ 2 ]

make_primes(L)

print S(N) + S(M)

print time.clock() - t