Diagonal Traverse

medium array matrix simulation

Problem

Given an m × n matrix, return all elements visited in diagonal zig-zag order: first diagonal from top-left moves up-right, the next moves down-left, and so on.

Inputmat = [[1,2,3],[4,5,6],[7,8,9]]

Output[1, 2, 4, 7, 5, 3, 6, 8, 9]

Cells on the same diagonal share the value r + c; we sweep diagonals in order, alternating direction.

matrix (rows by ;, cells by ,)

def find_diagonal_order(mat):
    m, n = len(mat), len(mat[0])
    out = []
    r = c = 0
    going_up = True
    for _ in range(m * n):
        out.append(mat[r][c])
        if going_up:
            if c == n - 1: r += 1; going_up = False
            elif r == 0: c += 1; going_up = False
            else: r -= 1; c += 1
        else:
            if r == m - 1: c += 1; going_up = True
            elif c == 0: r += 1; going_up = True
            else: r += 1; c -= 1
    return out

function findDiagonalOrder(mat) {
  const m = mat.length, n = mat[0].length;
  const out = [];
  let r = 0, c = 0, goingUp = true;
  for (let k = 0; k < m * n; k++) {
    out.push(mat[r][c]);
    if (goingUp) {
      if (c === n - 1) { r++; goingUp = false; }
      else if (r === 0) { c++; goingUp = false; }
      else { r--; c++; }
    } else {
      if (r === m - 1) { c++; goingUp = true; }
      else if (c === 0) { r++; goingUp = true; }
      else { r++; c--; }
    }
  }
  return out;
}

class Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        int m = mat.length, n = mat[0].length;
        int[] out = new int[m * n];
        int r = 0, c = 0; boolean goingUp = true;
        for (int k = 0; k < m * n; k++) {
            out[k] = mat[r][c];
            if (goingUp) {
                if (c == n - 1) { r++; goingUp = false; }
                else if (r == 0) { c++; goingUp = false; }
                else { r--; c++; }
            } else {
                if (r == m - 1) { c++; goingUp = true; }
                else if (c == 0) { r++; goingUp = true; }
                else { r++; c--; }
            }
        }
        return out;
    }
}

vector<int> findDiagonalOrder(vector<vector<int>>& mat) {
    int m = mat.size(), n = mat[0].size();
    vector<int> out(m * n);
    int r = 0, c = 0; bool goingUp = true;
    for (int k = 0; k < m * n; k++) {
        out[k] = mat[r][c];
        if (goingUp) {
            if (c == n - 1) { r++; goingUp = false; }
            else if (r == 0) { c++; goingUp = false; }
            else { r--; c++; }
        } else {
            if (r == m - 1) { c++; goingUp = true; }
            else if (c == 0) { r++; goingUp = true; }
            else { r++; c--; }
        }
    }
    return out;
}

Explanation

We need to read a grid in a zig-zag diagonal pattern: one diagonal goes up-and-right, the next goes down-and-left, and we keep flipping direction. The clean idea here is to track a single cursor (r, c) and a flag going_up that says which way we are currently moving.

On every step we record the value at mat[r][c], then decide where to go next based on the direction and whether we have hit an edge of the grid.

While moving up-right we normally do r -= 1; c += 1. But if we are about to fall off, we turn: hitting the right wall (c == n - 1) means step down (r += 1), and hitting the top wall (r == 0) means step right (c += 1). Each turn also flips going_up. The down-left case is the mirror image of this.

Example with [[1,2,3],[4,5,6],[7,8,9]]: we visit 1, go up-right but we are at the top so step right to 2, flip direction and go down-left to 4, hit the left wall so step down to 7, and so on, giving [1, 2, 4, 7, 5, 3, 6, 8, 9].

We visit every cell exactly once, so the total work is just the number of cells in the grid.

Time: O(m · n) Space: O(m · n) output