LeetCode 498. Diagonal Traverse

链接

https://leetcode-cn.com/problems/diagonal-traverse/description/

要求

给定一个含有 M x N 个元素的矩阵（M行，N列），请以对角线遍历的顺序返回这个矩阵中的所有元素，对角线遍历如下图所示。

输入:
[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]
输出:  [1,2,4,7,5,3,6,8,9]

思路

已写在代码中，效率较低。

代码

执行用时：912 ms

class Solution:
    def findDiagonalOrder(self, matrix):
        # 解题思路
        #-------------------------------------------------------
        # 1 2 3       1 2 3 x x
        # 4 5 6  ==>  x 4 5 6 x
        # 7 8 9       x x 7 8 9
        #-------------------------------------------------------

        # 特殊情况
        if not matrix:
            return matrix
        elif len(matrix) == 1:
            return matrix[0]
        elif len(matrix[0]) == 1:
            return [y for x in matrix for y in x]

        # 前后插入英文字母
        for matrix_index, matrix_list in enumerate(matrix):
            if matrix_index:
                matrix[matrix_index] = list(matrix_index * 'x') + matrix[matrix_index]
            matrix_counter = len(matrix) - matrix_index - 1
            if matrix_counter:
                matrix[matrix_index] = matrix[matrix_index] + list(matrix_counter * 'x')

        # zip函数解包重组
        matrix_zip = list(zip(*matrix))

        # 调整顺序，偶数位需要逆序
        for matrix_zip_index, matrix_zip_list in enumerate(matrix_zip):
            if matrix_zip_index % 2 == 0:
                matrix_zip[matrix_zip_index] = matrix_zip[matrix_zip_index][::-1]

        # 去除英文字符并重组为一个列表
        matrix_final_output = [y for x in matrix_zip for y in x if y != 'x']

        return matrix_final_output