Square Walker

#!/usr/bin/env python

import numpy as np
import argparse

def create_test(size, shuffle=False):
    data = np.arange(1, size*size+1)
    if shuffle:
        np.random.shuffle(data)
    return data.reshape(size, size)

def walk(data, size, level):
    steps=np.array([[0,1], [1,0], [0,-1], [-1,0]])
    walks=steps[np.arange(4).repeat(level*2)]
    start=(size - 1) / 2 - level
    return [data[tuple(x + (start, start))] for x in walks.cumsum(0)]

def offset(level):
    step=(level-1)*2+1
    off=[0, 1, step, 1, step, 1, step, 1]
    return np.cumsum(off)+(level-1)

def solve(data, size):
    result = [data[(size/2, size/2)]]
    rect1 = walk(data, size, 1)
    off = np.argmax(rect1)
    result = np.append(result, np.roll(rect1, -off))
    for level in range(2, size/2 + 1):
        rect = walk(data, size, level)
        tmp = np.roll(rect, -offset(level)[off])
        result = np.append(result, tmp)
    return result

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Square walker')
    parser.add_argument('size', help='size of data')
    parser.add_argument('-s', '--shuffle', action='store_true',
                        help='shuffle the test data')
    option = parser.parse_args()

    size = int(option.size)
    data = create_test(size, option.shuffle)
    print data
    print solve(data, size)