矩陣按對角線打印---python

將一個矩陣（二維數組）按對角線向右進行打印。（搜了一下發現好像是美團某次面試要求半小時手撕的題）
Example:
Input:
[
[1,2,3,4],
[5,1,2,3],
[9,5,1,2]
]
Output:
[[4],
[3, 3],
[2, 2, 2],
[1, 1, 1],
[5, 5],
[9]]

class Solution():
    def print_matix(self,list):
        print(list)
        rows = len(list)
        cols = len(list[0])
        result = []
        for k in range(2*cols-1):    # 共輸出 cols * 2 - 1 行
            diff = cols - k - 1  # 每一行的差
            for i in range(cols):    # 數組中每一個值的下標范圍是0到cols
                for j in range(cols):
                    if j-i == diff:
                        print(list[i][j])
                        result.append(list[i][j])
            print()

        print(result)
        return result

    def diagonal_right(self,matrix):
        if not matrix:
            return []
        row = len(matrix)
        col = len(matrix[0])
        col2 = col
        result = []
        for i in range(row):
            for j in range(col2 - 1, -1, -1):  # j倒序遍歷
                lst = []
                i1, j1 = i, j  # i1,j1用於方便同一對角線元素的添加，否則改變i,j影響開頭元素的選擇
                while i1 <= row - 1 and j1 <= col - 1:
                    lst.append(matrix[i1][j1])
                    j1 += 1
                    i1 += 1
                result.append(lst)
                if i == 0 and j == 0:  # 當遍歷完(0,0)開頭的一條對角線后，讓j固定為0
                    col2 = 1
        return ('the result is: %s' % result, 'end')



if __name__ == '__main__':
     list = [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]]
     solution = Solution()
     solution.print_matix(list)
     result = solution.diagonal_right(list)
     print(result)