Leetcode 498 Solution

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        vector<int> res;
        if (matrix.size() == 0 || matrix[0].size() == 0)
            return res;

        int m = matrix.size();
        int n = matrix[0].size();
        int i = 0, j = 0;
        int dir = 1;
        while (res.size() != m * n)
        {
            res.push_back(matrix[i][j]);

            int next_i = i;
            int next_j = j;

            if (dir == 1)
                next_i--, next_j++;
            else if (dir == 2)
                next_i++, next_j--;

            if (0 <= next_i && next_i < m && 0 <= next_j && next_j < n)
            {
                i = next_i;
                j = next_j;
            }
            else if (next_i >= m)
            {
                i = m - 1;
                j = next_j + 2;
                dir = 3 - dir;
            }
            else if (next_j >= n)
            {
                i = next_i + 2;
                j = n - 1;
                dir = 3 - dir;
            }
            else if (next_j < 0)
            {
                i = next_i;
                j = 0;
                dir = 3 - dir;
            }
            else if (next_i < 0)
            {
                i = 0;
                j = next_j;
                dir = 3 - dir;
            }
        }

        return res;
    }
};