Project Euler 128

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

Q128.
自然数を1から正六角形状のスパイラルに並べる。このとき、配置された数nと隣の6つの数との差のうちの素数の数をPD(n)とする。PD(n)=3となる2000番目のnは?

座標を、上にx、そこから左60°にyと定める。
最初、mapに数列を納めていったが、メモリが足りなかった。少し面倒だが、座標⇔数変換関数を作った。
それでも遅いので、図をよく見ていると、座標が(x, 0), (x, -1) (x ≥ 0)の数しかPD(n)=3になりえないことがわかった。



from math import sqrt
from itertools import count

nei = ( (-1, 1), (-1, 0), (0, -1), (1, -1), (1, 0), (0, 1) )
vec = ( (1, 0), (0, 1), (-1, 1), (-1, 0), (0, -1), (1, -1) )

def is_prime(n):
if n < 4:
return n != 1
else:
if (n & 1) == 0:
return False
p = 3
while p * p <= n:
if n % p == 0:
return False
p += 2
return True

def num_to_pt(n):
if n == 1:
return (0, 0)

k = (3 + int(sqrt(12 * n - 15))) / 6
m = n - 3 * k * (k - 1) - 2
r = m % k
m = m / k
return (vec[m][0] * k + nei[m][0] * r, vec[m][1] * k + nei[m][1] * r)

def pt_to_num(x, y):
if x == 0 and y == 0:
return 1

if y >= 0:
if x >= 0:
k = x + y
m = 0
r = y
elif x + y > 0:
k = y
m = 1
r = -x
else:
k = -x
m = 2
r = -x - y
else:
if x < 0:
k = -x - y
m = 3
r = -y
elif x + y < 0:
k = -y
m = 4
r = x
else:
k = x
m = 5
r = x + y

return 3 * k * (k - 1) + 2 + k * m + r

def PD(x, y):
n = pt_to_num(x, y)
counter = 0
for dx, dy in nei:
if is_prime(abs(pt_to_num(x + dx, y + dy) - n)):
counter += 1
return counter

N = 2000
counter = 0
for x in count():
if PD(x, 0) == 3:
counter += 1
if counter == N:
print pt_to_num(x, 0)
break
if PD(x, -1) == 3:
counter += 1
if counter == N:
print pt_to_num(x, -1)
break