python 实现zigzag排列

最近在写zigzag排列的时候，感觉还是有点难写的，主要是判断条件有点多。所以就放到这里做一下笔记。

def zigzag(data):  # 输入应该是np.array类型的数组
    (r, c) = data.shape
    if(r != c):
        print("行和列应该一样！")
        return 0
    ZZ = zeros((1, r*c))
    p = 0
    for index in range(2 * r):
        if(index <= r-1):
            for i in range(index+1):
                R = i
                C = index - i
                ZZ[0, p] = data[R, C]
                p = p + 1
        if(index > r-1):
            for i in range(2*r - index-1):
                C = r-1 - i
                R = index - C
                ZZ[0, p] = data[R, C]
                p += 1
    return ZZ

输入：

a = array([[1, 2, 4],
           [3, 5, 7],
           [6, 8, 9]])
b = array([[1, 2],
           [3, 4]])
c = array([[1, 2, 4, 7],
           [3, 5, 8, 11],
           [6, 9, 12, 14],
           [10,13, 15, 16]])

输出：