[LeetCode/Python/Java]Array and String - 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.

 

2차원 배열을 입력받았을 때 아래와 같은 순서로 배열을 출력한다.

코드를 보시려면 더보기를 클릭해 주세요.

파이썬 코드

class Solution:
    def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        if len(matrix) == 0:
            return []
        N, M = len(matrix), len(matrix[0])
        result, inter = [], []
        for d in range(N + M - 1):
            inter.clear()
            if d < M:
                r, c = 0, d
            else:
                r, c = d - M + 1, M - 1
            while r < N and c > -1:
                inter.append(matrix[r][c])
                r += 1
                c -= 1
            if d % 2 == 0:
                result.extend(inter[::-1])
            else:
                result.extend(inter)
        return result

 

자바 코드

자바 코드는 Solution에 있는 풀이 부분을 가져온 것입니다.

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0) {
            return new int[0];
        }
        int N = matrix.length;
        int M = matrix[0].length;
        int[] result = new int[N*M];
        int k = 0;
        ArrayList<Integer> intermediate = new ArrayList<Integer>();
        for (int d = 0; d < N + M - 1; d++) {
            intermediate.clear();
            int r = d < M ? 0 : d - M + 1;
            int c = d < M ? d : M - 1;
            while (r < N && c > -1) {
                
                intermediate.add(matrix[r][c]);
                ++r;
                --c;
            }
                
            if (d % 2 == 0) {
                Collections.reverse(intermediate);
            }
            
            for (int i = 0; i < intermediate.size(); i++) {
                result[k++] = intermediate.get(i);
            }
        }
        return result;
    }
}