498. Diagonal Traverse

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]]
Output: [1,2,4,7,5,3,6,8,9]

Example 2:

Input: mat = [[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

 

Solutions

Solution: Simulation

• Time complexity: O(mn)
• Space complexity: O(mn)

 

JavaScript

/**
 * @param {number[][]} mat
 * @return {number[]}
 */
const findDiagonalOrder = function (mat) {
  const m = mat.length;
  const n = mat[0].length;
  const result = [];
  let row = 0;
  let col = 0;
  let direction = 1;

  for (let index = 0; index < m * n; index++) {
    let nextRow = row - direction;
    let nextCol = col + direction;

    if (nextCol < 0) {
      nextCol = nextRow > m - 1 ? 1 : 0;
      nextRow = Math.min(m - 1, nextRow);
      direction *= -1;
    } else if (nextCol >= n) {
      nextCol = n - 1;
      nextRow = row + 1;
      direction *= -1;
    } else if (nextRow < 0) {
      nextRow = 0;
      direction *= -1;
    } else if (nextRow >= m) {
      nextRow = m - 1;
      nextCol = col + 1;
      direction *= -1;
    }

    result[index] = mat[row][col];
    row = nextRow;
    col = nextCol;
  }

  return result;
};