498. 对角线遍历

102 阅读1分钟

题目:
给你一个大小为 m x n 的矩阵 mat ,请以对角线遍历的顺序,用一个数组返回这个矩阵中的所有元素。 算法:
方法一:模拟
注意边界条件
1.矩阵左下角和右上角要特殊处理。
2.i >= n,j >= m的时候,左边移动方式不同于i < 0, j < 0 。

func findDiagonalOrder(mat [][]int) []int {
    n := len(mat)
    m := len(mat[0])
    count := m * n
    ans := make([]int, 0)
    dricts := [][]int{[]int{-1, 1}, []int{1, -1}}
    i, j := 0, 0
    flag := 0
    for count > 0 {
        count --
        // fmt.Println(ans, i, j)
        ans = append(ans, mat[i][j])
        i, j = i + dricts[flag][0],  j + dricts[flag][1]
        if 0 <= i && i < n && 0 <= j && j < m {
            continue
        }
        if i == -1 && j == m {
            i = i + 2
            j = j - 1
        }

        if i == n && j == -1 {
            i = i - 1
            j = j + 2
        }
        if i < 0 {
            i ++
        }
        if j < 0 {
            j ++
        }
        if i >= n {
            i -- 
            j = j + 2
        }
        if j >= m {
            i = i + 2
            j --
        }
        
        flag = flag ^ 1
    }
    return ans
}

方法二:更简洁的模拟

func findDiagonalOrder(mat [][]int) []int {
    n := len(mat)
    m := len(mat[0])
    count := m * n
    ans := make([]int, 0)
    dricts := [][]int{[]int{-1, 1}, []int{1, -1}}
    i, j := 0, 0
    flag := 0
    for count > 0 {
        count --
        // fmt.Println(ans, i, j)
        ans = append(ans, mat[i][j])
        x, y := i + dricts[flag][0],  j + dricts[flag][1]
        if x < 0  || x >= n || y < 0 || y >= m {
            if flag == 0 {
                if j + 1 < m {
                    x = i
                    y = j + 1
                } else {
                    x = i + 1
                    y = j
                }
            } else {
                if i + 1 < n {
                    x = i + 1
                    y = j
                } else {
                    x = i
                    y = j + 1
                }
            }
            flag = flag ^ 1
        }
        
        i = x
        j = y
        
    }
    return ans
}