一个蛇形矩阵问题

一个蛇形矩阵问题

话不多说，先上代码：

# -*- coding: utf-8 -*-
"""
Created on Tue Aug 21 18:12:36 2018

@author: duanshao
"""

import numpy as np

def snake(n):
    raw = [0] * n
    rs = []
    for i in range(n): rs.append(raw.copy())

    i, j = 0 , 0
    direct = True
    for num in range(1, n*n+1):
        rs[i][j] = num

        if direct:
            j += 1
            i -= 1
        else:
            j -= 1
            i += 1
        ci, cj = i, j

        if i < 0 or j < 0 or i > n-1 or j > n-1:
            direct = not direct

        if i > n-1:
            ci = n-1
            cj += 2
        elif i < 0 and j < n:
            ci = 0

        if j > n-1:
            cj = n-1                
            ci += 2
        elif j < 0 and i < n:
            cj = 0
        i, j = ci, cj
    print(np.array(rs))

snake(5)

结果测试时需要注意函数输入为1，2，奇数和偶数情况时，程序是否有误。
测试结果：
snake(5):

snake(4):