498_Diagonal Traverse

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]

Solution 1: brute force

Idea:

• Have a flag to show if the direction goes up right or down left, take different operation for different direction, update the flag when approach the border.

• Take special care of elements on the border of the matrix, and the update rules are different for the four borders.

Time Complexity: $O(mn)$

Space Complexity: $O(mn)$

def findDiagonalOrder(matrix):
    """
    :type matrix: List[List[int]]
    :rtype: List[int]
    """
    if not matrix: return []
    output = []
    m = len(matrix)
    n = len(matrix[0])
    i, j, flag = 0, 0, 1
    while i < m and j < n:
        output.append(matrix[i][j])
        if flag == 1: # go up right
            if j == n - 1:
                i += 1
                flag *= -1
            elif i == 0:
                j += 1
                flag *= -1
            else:
                i -= 1
                j += 1
        else: # go down left
            if i == m - 1:
                j += 1
                flag *= -1
            elif j == 0:
                i += 1
                flag *= -1
            else:
                i += 1
                j -= 1       
    return output

Solution 2: use sorting

Idea:

• In each diagonal, col index + row index is a constant. We are sorting this constant in ascending order.

• In a diagonal, the direction which loops each element is decided by if col index + row index is odd or even. If it's even, it traverse larger rows, smaller cols first. If it's odd, it traverse larger cols, small rows first.

Time Complexity: $O(mn \log{mn})$

Space Complexity: $O(mn)$

def findDiagonalOrder(matrix):
    """
    :type matrix: List[List[int]]
    :rtype: List[int]
    """
    l = [(i, j) for i in range(len(matrix)) for j in range(len(matrix[0]))]
    l.sort(key=lambda x: sum(x) * (len(matrix) + len(matrix[0]) ) - x[sum(x)%2])
    return [matrix[x][y] for x, y in l]