[Leetcode] 498. Diagonal Traverse

Well, this problem itself has nothing to note. But I must take a screen shot for record - the first time my program beats 100%... just for fun. Problem wise, it's just a bunch of index tricks - along the diagonal, element matrix[i][j], i + j maintains the same.

---

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:

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

Note:

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

---

One thing to note is that the input size says that we can only afford $O(n)$ algorithm.

class Solution(object):
    def findDiagonalOrder(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: List[int]
        """
        if not any(matrix): return []
        m, n = len(matrix), len(matrix[0])
        
        if m == 1: return matrix[0]
        if n == 1: return [matrix[i][0] for i in range(m)]
        
        res = []
        for k in range(m+n):
            if k >= n:
                start = k -n + 1 # row 
            else:
                start = 0
            if k >= m:
                end = m - 1 
            else:
                end = k
            tmp = []
            for i in range(start, end + 1):
                tmp.append(matrix[i][k-i])
            
            if k % 2 == 0:
                res += tmp[::-1]
            else:
                res += tmp
        return res