62 leetcode - Spiral Matrix II

#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
Spiral Matrix II
Given an integer n, generate a square matrix filled with elements from 1 to n2 in spiral order.
For example,
Given n = 3,
You should return the following matrix:
[
[ 1, 2, 3 ],
[ 8, 9, 4 ],
[ 7, 6, 5 ]
]
'''
class Solution(object):
def generateMatrix(self, n):
"""
:type n: int
:rtype: List[List[int]]
"""
if n == 0:
return []
matrix = [[None] * n for i in range(n)]
up = left = 0
down = right = n - 1
nums = 1
while down >= up and left <= right:
col = left
while col <= right:
matrix[up][col] = nums
nums += 1
col += 1
up += 1

row = up
while row <= down:
matrix[row][right] = nums
nums += 1
row += 1
right -= 1

col = right
while col >= left:
matrix[down][col] = nums
nums += 1
col -= 1
down -= 1

row = down
while row >= up:
matrix[row][left] = nums
nums += 1
row -= 1
left += 1

return matrix

if __name__ == "__main__":
s = Solution()
print s.generateMatrix(1)
print s.generateMatrix(2)
print s.generateMatrix(3)
print s.generateMatrix(4)