Diagonal Traverse

Leet Code Link - https://leetcode.com/problems/diagonal-traverse/description/

Complexity - Medium

Description

Given an m x n matrix mat, return an array of all the elements of the array in a diagonal order.

Example 1:

Input: mat = [1,2,3],[4,5,6],[7,8,9](/shilpathota/99-leetcode-solutions/wiki/1,2,3],[4,5,6],[7,8,9)
Output: [1,2,4,7,5,3,6,8,9]

Example 2:

Input: mat = [1,2],[3,4](/shilpathota/99-leetcode-solutions/wiki/1,2],[3,4)
Output: [1,2,3,4]

Constraints:

m == mat.length
n == mat[i].length
1 <= m, n <= 104
1 <= m * n <= 104
-105 <= mat[i][j] <= 105

Solution

It's pure simulation. However, in order to implement this simulation, we need to understand the walking patterns inside the array. Basically, in the previous approach, figuring out the head of the diagonal was pretty easy. In this case, it won't be that easy. We need to figure out two things for each diagonal:

The direction in which we want to process it's elements and The head or the starting point for the diagonal depending upon its direction. Let's see these two things annotated on a sample matrix.

Now that we know what two things we need to figure out, let's get to the part where we actually do it! The direction is pretty straightforward. We can simply use a boolean variable and keep alternating it to figure out the direction for a diagonal. That part is sorted. The slightly tricky part is figuring out the head of the next diagonal.

The good part is, we already know the end of the previous diagonal. We can use that information to figure out the head of the next diagonal.

Next head when going UP

Let's look at the two scenarios that we may come across when we are at the tail end of a downwards diagonal and we want to find the head of the next diagonal.

So, the general rule that we will be following when we want to find the head for an upwards going diagonal is that:

The head would be the node directly below the tail of the previous diagonal. Unless the tail lies in the last row of the matrix in which case the head would be the node right next to the tail.

Next head when going DOWN

Let's look at the two scenarios that we may come across when we are at the tail end of an upwards diagonal and we want to find the head of the next diagonal.

So, the general rule that we will be following when we want to find the head for a downwards going diagonal is that:

The head would be the node to the right of the tail of the previous diagonal. Unless the tail lies in the last column of the matrix in which case the head would be the node directly below the tail.

Algorithm

Initialize a boolean variable called direction which will tell us whether the current diagonal is an upwards or downwards going. Based on the current direction and the tail, we will determine the head of the next diagonal. Initially the direction would be 1 which would indicate up. We will keep alternating this value from one iteration to the next.

Assuming we know the head of a diagonal, say matrix[i][j], we will use the direction to progress along the diagonal and process its elements.

For an upwards going diagonal, the next element in the diagonal would be matrix[i−1][j+1] For a downwards going diagonal, the next element would be matrix[i+1][j−1]. We keep processing the elements of the current diagonal until we go out of the boundaries of the matrix.

Now, given that we know the tail of the diagonal (the last node before we went out of bounds), let's see how we can find the next head. Note that in the following pseudocode, the direction is for the current diagonal and we are trying to find the head of the next diagonal. So, if the direction is up, it means the next diagonal would be going down and vice-versa.

tail = [i, j]
if direction == up, then {
   if [i, j + 1] is within bounds, then {
       next_head = [i, j + 1]
   } else { 
       next_head = [i + 1, j]
   }
} else {
   if [i + 1, j] is within bounds, then {
       next_head = [i + 1, j]
   } else { 
       next_head = [i, j + 1]
   }
}

We keep processing the elements of a diagonal and once the current diagonal ends, we use the current direction and the tail element to find the next head and we switch over to processing the next diagonal. Also remember to flip the direction bit.

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        if(mat == null || mat.length == 0){
            return new int[0];
        }

        int N = mat.length;
        int M = mat[0].length;

        int row= 0,col =0;

        int direction =1;

        int[] result = new int[N*M];
        int r=0;
        while(row<N && col < M){
            result[r++] = mat[row][col];

            int new_row = row + (direction == 1 ?-1:1);
            int new_col = col + (direction == 1 ?1:-1);
            if(new_row<0||new_col<0 || new_row==N||new_col==M){
                if(direction == 1){
                    row+=(col == M-1?1:0);
                    col+=(col < M-1?1:0);
                }
                else{
                    col += (row==N-1?1:0);
                    row += (row<N-1?1:0);
                }
                direction = 1-direction;
            }else{
                row = new_row;
                col = new_col;
            }
        }
                    return result;

    }
}

Complexity

Time Complexity: O(N⋅M) since we process each element of the matrix exactly once.

Space Complexity: O(1) since we don't make use of any additional data structure. Note that the space occupied by the output array doesn't count towards the space complexity since that is a requirement of the problem itself. Space complexity comprises any additional space that we may have used to get to build the final array. For the previous solution, it was the intermediate arrays. In this solution, we don't have any additional space apart from a couple of variables