# ulam spiral, rotate 45 degree

primes_ulam_spiral_rotated45.py

# ulam spiral, rotate 45 degree
# http://users.livejournal.com/_winnie/392996.html

import Image

def is_prime(n):
    if n <= 2:
        return 1
    k = 2
    while k*k <= n:
        if n % k == 0:
            return 0
        k += 1
    return 1

A = set(z*z - z + 41  for z in xrange(20000))
B = set(2*z*z + 40*z + 1 for z in xrange(20000))

def ulam_spiral(n):
    x, y = 1, 0
    for k in xrange(1, n):
        #print x, y
        for _j in xrange(2*k-1):
            yield (k, x, y)
            y -= 1
        for _j in xrange(2*k):
            yield (k, x, y)
            x -= 1
        for _j in xrange(2*k):
            yield (k, x, y)
            y += 1
        for _j in xrange(2*k+1):
            yield (k, x, y)
            x += 1

def text_print():
    s = { }
    i = 1
    for _k, x, y in ulam_spiral(3):
        i += 1
        print x, y
        assert (x, y) not in s
        s[(x, y)] = i
    xs = [x for x, y in s.keys()]
    ys = [y for x, y in s.keys()]
    for y in xrange(min(ys), max(ys)+1):
        for x in xrange(min(xs), max(xs)+1):
            if (x, y) in s:
                print '% 2s' % s[(x, y)],
            else:
                print '  ',
        print

def main():
    i = 1
    n = 1000
    img = Image.new("RGB", (n*2 + 1, n*2 + 1))
    for _k, x, y in ulam_spiral(n):
        i += 1
        c = is_prime(i)
        r, g, b = c*255, c*255, c*255
        if c and (i in A):
            r, g, b = r, 0, 0
        if r or g or b:
            x2 = (x+y)/2 #rotate 45 degree
            y2 = (x-y)/2
            img.putpixel((x2 + n, y2 + n), (r, g, b))
    img.show()

main()