【算法练习】498. 对角线遍历（模拟）

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;
        int[] ans = new int[m * n];
        int cnt = 0;
        int dir = 1;
        int i = 0, j = 0;
        while (cnt != m * n) {
            ans[cnt++] = mat[i][j];
            if (cnt >= m * n) break;
            int ti = i, tj = j;
            if (dir == 1) {
                // 斜向上遍历
                ti = ti - 1;
                tj = tj + 1;
            } else {
                // 斜向下遍历
                ti = ti + 1;
                tj = tj - 1;
            }
            if (ti < 0 || ti >= m || tj < 0 || tj >= n) {
                // 越界处理
                if (dir == 1) {
                    ti = j + 1 < n ? i : i + 1;
                    tj = j + 1 < n ? j + 1 : j;
                } else {
                    ti = i + 1 < m ? i + 1 : i;
                    tj = i + 1 < m ? j : j + 1;
                }
                dir *= -1;  // 只有越界时需要转换方向
            }
            i = ti;
            j = tj;
        }
        return ans;
    }
}