498. Diagonal Traverse

Description

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:

image

Note:

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

Solution

Enum index sum, O(mn), S(mn)

利用规律：对于每条对角线来说，所有节点的i + j都相等。
所以可以枚举sum = i + j，然后一条一条打印即可。

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
            return new int[0];
        }
        
        List<Integer> list = new ArrayList<>();
        int m = matrix.length;
        int n = matrix[0].length;
        for (int sum = 0; sum <= m + n - 2; ++sum) { // enum sum of rowIndex and colIndex
            List<Integer> diagonal = new ArrayList<>();
            
            for (int i = 0; i <= Math.min(m - 1, sum); ++i) {
                if (sum - i < n) {
                    diagonal.add(matrix[i][sum - i]);
                }
            }
            
            if (sum % 2 == 0) { // reverse if needed
                Collections.reverse(diagonal);
            }
            list.addAll(diagonal);
        }
        
        return list.stream().mapToInt(i->i).toArray();
    }
}

Enum index sum, O(mn), S(1)

基于上面的算法，可以省略掉list，直接向int[] res中写。

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
            return new int[0];
        }
        
        int m = matrix.length;
        int n = matrix[0].length;
        int[] res = new int[m * n];
        int count = 0;
        
        for (int sum = 0; sum <= m + n - 2; ++sum) { // enum sum of rowIndex and colIndex
            if (sum % 2 != 0) {
                for (int i = 0; i <= Math.min(m - 1, sum); ++i) {
                    if (sum - i < n) {
                        res[count++] = matrix[i][sum - i];
                    }
                }
            } else {    // reverse
                for (int i = Math.min(m - 1, sum); i >= 0; --i) {
                    if (sum - i < n) {
                        res[count++] = matrix[i][sum - i];
                    }
                }
            }
        }
        
        return res;
    }
}