题目要求

思路：模拟

• 定义一个$dir$表示两个方向，根据方向的不同向下一个移动，遇到需要拐弯的时候根据越界情况判断。

Java

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        int n = mat.length, m = mat[0].length, cnt = m * n;
        int[] res = new int[cnt];
        int x = 0, y = 0, dir = 1, idx = 0;
        while(idx != cnt) {
            res[idx++] = mat[x][y];
            int nx = x, ny = y;
            if(dir == 1) { // 左下
                nx = x - 1;
                ny = y + 1;
            }
            else { // 右上
                nx = x + 1;
                ny = y - 1;
            }
            if(idx < cnt && (nx < 0 || nx >= n || ny < 0 || ny >= m)) { // 拐弯
                if(dir == 1) {
                    // 纵向越界向右，否则向下
                    nx = y + 1 < m ? x : x + 1;
                    ny = y + 1 < m ? y + 1 : y;
                }
                else {
                    // 横向越界向下，否则向右
                    nx = x + 1 < n ? x + 1 : x;
                    ny = x + 1 < n ? y : y + 1;
                }
                dir *= -1;
            }
            x = nx;
            y = ny;
        }
        return res;
    }
}

• 时间复杂度：$O(m\times n)$
• 空间复杂度：$O(1)$

C++

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
        int n = mat.size(), m = mat[0].size(), cnt = m * n;
        vector<int> res(cnt);
        int x = 0, y = 0, dir = 1, idx = 0;
        while(idx != cnt) {
            res[idx++] = mat[x][y];
            int nx = x, ny = y;
            if(dir == 1) { // 左下
                nx = x - 1;
                ny = y + 1;
            }
            else { // 右上
                nx = x + 1;
                ny = y - 1;
            }
            if(idx < cnt && (nx < 0 || nx >= n || ny < 0 || ny >= m)) { // 拐弯
                if(dir == 1) {
                    // 纵向越界向右，否则向下
                    nx = y + 1 < m ? x : x + 1;
                    ny = y + 1 < m ? y + 1 : y;
                }
                else {
                    // 横向越界向下，否则向右
                    nx = x + 1 < n ? x + 1 : x;
                    ny = x + 1 < n ? y : y + 1;
                }
                dir *= -1;
            }
            x = nx;
            y = ny;
        }
        return res;
    }
};

• 时间复杂度：$O(m\times n)$
• 空间复杂度：$O(1)$

Rust

impl Solution {
    pub fn find_diagonal_order(mat: Vec<Vec<i32>>) -> Vec<i32> {
        let n = mat.len();
        let m = mat[0].len();
        let cnt = m * n;
        let mut res = vec![0; cnt];
        let (mut x, mut y, mut dir, mut idx) = (0, 0, 1, 0);
        while idx != cnt {
            res[idx] = mat[x][y];
            idx += 1;
            let (mut nx, mut ny) = (x, y);
            if dir == 1 { // 左下
                nx = x - 1;
                ny = y + 1;
            }
            else { // 右上
                nx = x + 1;
                ny = y - 1;
            }
            if idx < cnt && (nx < 0 || nx >= n || ny < 0 || ny >= m) {
                if dir == 1 {
                    // 纵向越界向右，否则向下
                    if y + 1 < m {
                        nx = x;
                        ny = y + 1;
                    }
                    else {
                        nx = x + 1;
                        ny = y;
                    }
                }
                else {
                    // 横向越界向下，否则向右
                    if x + 1 < n {
                        nx = x + 1;
                        ny = y;
                    }
                    else {
                        nx = x;
                        ny = y + 1;
                    }
                }
                dir *= -1;
            }
            x = nx;
            y = ny;
        }
        res
    }
}

• 时间复杂度：$O(m\times n)$
• 空间复杂度：$O(1)$

总结

直接模拟、完全不是中等题难度……甚至可以直接迁移到Rust里。