Diagonal Traverse — Leetcode 498

Original problem link: https://leetcode.com/problems/diagonal-traverse/description/

In this problem, we’re given an mxn matrix and asked to return a 1D array of its values in diagonal order.

I noticed a pattern where for each diagonal row, the sum of the indices would be the same and we kept increasing this sum for each new diagonal. Thus, I used this sum to simulate the diagonals.

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        int[] result = new int[mat.length * mat[0].length];
        boolean isUp = true;
        int next = 0;
        int i = 0;
        int j = 0;
        int desiredSum = 0;

        while (next < result.length) {
            while (next < result.length && i >= 0 && j >= 0 && i < mat.length && j < mat[0].length) {
                result[next] = mat[i][j];
                if (isUp) {
                    i--;
                    j++;
                } else {
                    i++;
                    j--;
                }
                next++;
            }

            desiredSum++;
            isUp = !isUp;
            if (i < 0) {
                i = 0;
                j = desiredSum;
            }
            if (j < 0) {
                j = 0;
                i = desiredSum;
            }
            if (i == mat.length) {
                i = mat.length - 1;
                j = desiredSum - mat.length + 1;
            }
            if (j == mat[0].length) {
                i = desiredSum - mat[0].length + 1;
                j = mat[0].length - 1;
            }
        }

        return result;
    }
}

The runtime is between 2–3 ms (beating 93%-49% of submissions). It’s faster or the same as the Leetcode official solutions :)