DiagonalTraverse498.java

/**
 * 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.
 * 
 * Example:
 * Input:
 * [
 * [ 1, 2, 3 ],
 * [ 4, 5, 6 ],
 * [ 7, 8, 9 ]
 * ]
 * Output: [1,2,4,7,5,3,6,8,9]
 * 
 * Explanation:
 * https://leetcode.com/static/images/problemset/diagonal_traverse.png
 * 
 * Note:
 * The total number of elements of the given matrix will not exceed 10,000.
 */


public class DiagonalTraverse498 {
 public int[] findDiagonalOrder(int[][] matrix) {
 if (matrix == null || matrix.length == 0 || matrix[0].length == 0) return new int[0];
 int M = matrix.length;
 int N = matrix[0].length;
 int total = M * N;
 int[] res = new int[total];
 int i = 0;
 int[] pos = new int[2];
 res[i++] = matrix[0][0];
 while (i < total) {
 pos = next(pos, M, N);
 res[i++] = matrix[pos[0]][pos[1]];
 }
 return res;
 }

 private int[] next(int[] pos, int M, int N) {
 int i = pos[0];
 int j = pos[1];
 if (isEven(i+j)) { // top-right
 int ni = i - 1;
 int nj = j + 1;
 if (nj == N) {
 nj = N - 1;
 ni = i + 1;
 } else if (ni < 0) {
 ni = 0;
 }
 return new int[]{ni, nj};
 } else { // bottom-left
 int ni = i + 1;
 int nj = j - 1;
 if (ni == M) {
 ni = M - 1;
 nj = j + 1;
 } else if (nj < 0) {
 nj = 0;
 }
 return new int[]{ni, nj};
 }
 }
 
 private boolean isEven(int x) {
 return x % 2 == 0;
 }

}