Diagonal Traverse LeetCode solution in Java

1
/**
2
 * Simulate Diagonal Order Traversal
3
 *
4
 * <p>r+c determines which diagonal you are on. For ex: [2,0],[1,1],[0,2] are all on same diagonal
5
 * with r+c =2. If you check the directions of diagonals, first diagonal is up, second diagonal is
6
 * down, third one is up and so on.. Therefore (r+c)%2 simply determines direction. Even is UP
7
 * direction. Odd is DOWN direction.
8
 *
9
 * <p>Time Complexity: O(M*N)
10
 *
11
 * <p>Space Complexity: O(1) without considering result space.
12
 *
13
 * <p>M = Number of rows. N = Number of columns.
14
 */
15
class Solution {
16
  public int[] findDiagonalOrder(int[][] matrix) {
17
    if (matrix == null) {
18
      throw new IllegalArgumentException("Input matrix is null");
19
    }
20
    if (matrix.length == 0 || matrix[0].length == 0) {
21
      return new int[0];
22
    }
23

24
    int rows = matrix.length;
25
    int cols = matrix[0].length;
26
    int[] result = new int[rows * cols];
27
    int r = 0;
28
    int c = 0;
29

30
    for (int i = 0; i < result.length; i++) {
31
      result[i] = matrix[r][c];
32
      if ((r + c) % 2 == 0) { // Move Up
33
        if (c == cols - 1) {
34
          // Reached last column. Now move to below cell in the same column.
35
          // This condition needs to be checked first due to top right corner cell.
36
          r++;
37
        } else if (r == 0) {
38
          // Reached first row. Now move to next cell in the same row.
39
          c++;
40
        } else {
41
          // Somewhere in middle. Keep going up diagonally.
42
          r--;
43
          c++;
44
        }
45
      } else { // Move Down
46
        if (r == rows - 1) {
47
          // Reached last row. Now move to next cell in same row.
48
          // This condition needs to be checked first due to bottom left corner cell.
49
          c++;
50
        } else if (c == 0) {
51
          // Reached first columns. Now move to below cell in the same column.
52
          r++;
53
        } else {
54
          // Somewhere in middle. Keep going down diagonally.
55
          r++;
56
          c--;
57
        }
58
      }
59
    }
60

61
    return result;
62
  }
63
}