Problem Description

You are given a 2D matrix of size m x n (with m rows and n columns). Your task is to return all elements of the matrix in a diagonal order, starting from the top-left element. The diagonal order follows a zigzag pattern: the first diagonal goes up-right, the next down-left, and so on, alternating directions for each diagonal.

• Each element in the matrix must be visited exactly once and in the correct diagonal order.
• Do not revisit or skip any elements.
• The output should be a single list (or array) of all matrix elements in the required diagonal traversal order.

For example, given the matrix:

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

The output should be: [1,2,4,7,5,3,6,8,9]

Thought Process

The main challenge is to traverse the matrix in a zigzag diagonal pattern, alternating directions for each diagonal. A brute-force solution might attempt to keep track of visited elements or simulate the direction changes manually, but this quickly becomes complex and error-prone.

Observing the problem, we can note that each diagonal in the matrix can be identified by the sum of its row and column indices (i.e., for element matrix[i][j], the diagonal number is i + j). This insight helps us group elements by diagonals and then process each group in the correct order (reversing every other diagonal as needed).

Instead of using extra space to track visited elements, we can directly calculate the next position based on our current direction and the matrix boundaries, ensuring we visit each element once.

Solution Approach

Here is a step-by-step breakdown of the optimal solution:

1. Identify Diagonals: For each element at position (i, j), the sum i + j is constant for all elements in the same diagonal.
2. Group Elements: We can use a list (or map) to collect elements by their diagonal number (i + j).
3. Build Output: For each diagonal:
   • If the diagonal number is even, reverse the order of elements (since the traversal is up-right).
   • If the diagonal number is odd, keep the order as is (down-left).
4. Concatenate Results: Combine all diagonals in order to produce the final traversal.

This approach is both simple and efficient, avoiding the need for explicit direction simulation or extra visited markers.

Example Walkthrough

Consider the input matrix:

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

1. Identify Diagonals:
   • Diagonal 0: (0,0) → [1]
   • Diagonal 1: (0,1), (1,0) → [2,4]
   • Diagonal 2: (0,2), (1,1), (2,0) → [3,5,7]
   • Diagonal 3: (1,2), (2,1) → [6,8]
   • Diagonal 4: (2,2) → [9]
2. Process Each Diagonal:
   • Diagonal 0 (even): reverse [1] → [1]
   • Diagonal 1 (odd): keep as is → [2,4]
   • Diagonal 2 (even): reverse [3,5,7] → [7,5,3]
   • Diagonal 3 (odd): keep as is → [6,8]
   • Diagonal 4 (even): reverse [9] → [9]
3. Concatenate: [1] + [2,4] + [7,5,3] + [6,8] + [9] = [1,2,4,7,5,3,6,8,9]

Code Implementation

class Solution:
    def findDiagonalOrder(self, matrix):
        if not matrix or not matrix[0]:
            return []
        m, n = len(matrix), len(matrix[0])
        diagonals = {}
        for i in range(m):
            for j in range(n):
                if i + j not in diagonals:
                    diagonals[i + j] = []
                diagonals[i + j].append(matrix[i][j])
        result = []
        for k in range(m + n - 1):
            if k % 2 == 0:
                result.extend(reversed(diagonals[k]))
            else:
                result.extend(diagonals[k])
        return result
      

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        if (matrix.empty() || matrix[0].empty()) return {};
        int m = matrix.size(), n = matrix[0].size();
        vector<vector<int>> diagonals(m + n - 1);
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                diagonals[i + j].push_back(matrix[i][j]);
            }
        }
        vector<int> result;
        for (int k = 0; k < m + n - 1; ++k) {
            if (k % 2 == 0) {
                result.insert(result.end(), diagonals[k].rbegin(), diagonals[k].rend());
            } else {
                result.insert(result.end(), diagonals[k].begin(), diagonals[k].end());
            }
        }
        return result;
    }
};
      

class Solution {
    public int[] findDiagonalOrder(int[][] matrix) {
        if (matrix == null || matrix.length == 0) return new int[0];
        int m = matrix.length, n = matrix[0].length;
        List<List<Integer>> diagonals = new ArrayList<>();
        for (int i = 0; i < m + n - 1; ++i) diagonals.add(new ArrayList<>());
        for (int i = 0; i < m; ++i) {
            for (int j = 0; j < n; ++j) {
                diagonals.get(i + j).add(matrix[i][j]);
            }
        }
        int[] result = new int[m * n];
        int idx = 0;
        for (int k = 0; k < m + n - 1; ++k) {
            List<Integer> diagonal = diagonals.get(k);
            if (k % 2 == 0) {
                for (int i = diagonal.size() - 1; i >= 0; --i)
                    result[idx++] = diagonal.get(i);
            } else {
                for (int num : diagonal)
                    result[idx++] = num;
            }
        }
        return result;
    }
}
      

var findDiagonalOrder = function(matrix) {
    if (!matrix.length || !matrix[0].length) return [];
    const m = matrix.length, n = matrix[0].length;
    const diagonals = Array(m + n - 1).fill(0).map(() => []);
    for (let i = 0; i < m; ++i) {
        for (let j = 0; j < n; ++j) {
            diagonals[i + j].push(matrix[i][j]);
        }
    }
    const result = [];
    for (let k = 0; k < m + n - 1; ++k) {
        if (k % 2 === 0) {
            result.push(...diagonals[k].reverse());
        } else {
            result.push(...diagonals[k]);
        }
    }
    return result;
};
      

Time and Space Complexity

• Brute-force approach:
  • Could involve simulating movement and marking visited cells, leading to O(mn) time and O(mn) space.
• Optimized approach (as above):
  • Time complexity: O(mn), since every element is visited and added to the output exactly once.
  • Space complexity: O(mn), for storing the diagonals and the output array.

The solution is efficient and scales linearly with the size of the matrix.

Summary

The key insight for this problem is that elements on the same diagonal share the same sum of their indices. By grouping elements by this sum and reversing every other group, we can efficiently produce the required zigzag diagonal traversal. This approach is both simple and robust, avoiding complex direction simulation and ensuring every element is visited exactly once.

The result is an elegant, efficient, and easy-to-understand solution that leverages the properties of matrix indices.