59_Spiral Matrix II

Given a positive integer n , generate a square matrix filled with elements from 1 to n^2 in spiral order.

Example:

Input:
 3

Output:
[
 [ 1, 2, 3 ],
 [ 8, 9, 4 ],
 [ 7, 6, 5 ]
]

Solution:

The idea is the same as leetcode 54. Spiral Matrix

def generate_spiral_matrix(n):
    """
    :type n: int
    :rtype: List[List[int]]
    """
    output = [[0 for i in range(n)] for j in range(n)]
    shift = [[0,1], [1,0], [0,-1], [-1,0]]
    direction = 0
    i = j = 0
    for num in range(1, n**2+1) :
        output[i][j] = num
        i_new, j_new = i + shift[direction][0], j + shift[direction][1]
        if i_new not in range(n) or j_new not in range(n) or output[i_new][j_new] != 0 :
            direction = (direction + 1) % 4
            i_new, j_new = i + shift[direction][0], j + shift[direction][1]
        i, j = i_new, j_new
    return output