[LeetCode / Kotlin] Diagonal Traverse

문제

 

Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

 

Example 1:

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

 

 

Example 2:

Input: mat = [[1,2],[3,4]]
Output: [1,2,3,4]

 

 

Constrains:

• m == mat.length
• n == mat[i].length
• 1 <= m, n <= 104
• 1 <= m * n <= 104
• -105 <= mat[i][j] <= 105

 

나의 풀이

각 이동방향의 좌표 변화를 정의하고, 충돌케이스를 정의하였다.

class Solution {
    fun findDiagonalOrder(mat: Array<IntArray>): IntArray {
        val maxM = mat.lastIndex
        if (maxM < 0) {
            return IntArray(0)
        }
        
        val maxN = mat[0].lastIndex
        if (maxN < 0) {
            return IntArray(0)
        }
        
        // true가 위 false가 아래
        var direction = true
        var m = 0
        var n = 0
        val res = IntArray((maxM+1) * (maxN+1)).apply {
            this[0] = mat[m][n]
        }
        
        
        for (idx in 1..res.lastIndex) {
            
            if (direction) {
                if (m - 1 >= 0 && n + 1 <= maxN) {
                    // normal case
                    m -= 1
                    n += 1
                } else {
                    direction = !direction
                    if (m - 1 < 0 && n + 1 <= maxN) {
                        n += 1
                    } else {
                        m += 1
                    }
                }
                
            } else {
                if (m + 1 <= maxM && n - 1 >= 0) {
                    // normal case
                    m += 1
                    n -= 1
                } else {
                    direction = !direction
                    if (m + 1 <= maxM && n -1 < 0) {
                        m += 1
                    } else {
                        n += 1
                    }
                }
            }
            
            res[idx] = mat[m][n]
        }
        
        return res
    }
    
}