Given a matrix of M x N elements (M rows, N columns), return all elements of the matrix in diagonal order as shown in the below image.

Example:

1
2
3
4
5
6
7
Input:
[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]
Output:  [1,2,4,7,5,3,6,8,9]

Note:

  1. The total number of elements of the given matrix will not exceed 10,000.

思路

直接按照题目给的顺序来进行访问，时间复杂度是O(n)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class Solution(object):
    def findDiagonalOrder(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: List[int]
        """
        x = 0
        y = 1
        res = []
        vertical = len(matrix)
        horizontal = 0
        if (vertical):
            horizontal = len(matrix[0])
        
        if (horizontal > 1 and vertical > 1):
            res.append(matrix[0][0])
            forward = False
            downward = True
            upward = False
            count = 0
            while( x < vertical - 1 or y < horizontal - 1 ):
                res.append(matrix[x][y])
                
                if (forward):
                    if ( x < vertical-1 and y == 0 ):
                        x += 1
                        upward = True
                        forward = False
                    elif (x < vertical-1 and y == horizontal - 1):
                        x += 1
                        downward = True
                        forward = False
                    else:
                        if (x == 0):
                            y += 1
                            downward = True
                            forward = False
                        elif (x == vertical -1):
                            y += 1
                            upward = True
                            forward = False
                else:
                    if (upward):
                        x -= 1
                        y += 1
                        if (x == 0 or y == horizontal - 1):
                            upward = False
                            forward = True
                    
                    if (downward):
                        x += 1
                        y -= 1
                        if (x == vertical - 1 or y == 0):
                            downward = False
                            forward = True
            
            if (x == vertical - 1 and y == horizontal - 1):
                res.append(matrix[x][y])
        
        else:
            if (horizontal == 1):
                for i in matrix:
                    res.append(i[0])
            elif(vertical == 1):
                for i in matrix[0]:
                    res.append(i)
        return res