[leetcode]498. Diagonal Traverse

题目链接：https://leetcode.com/problems/diagonal-traverse/?tab=Description

Given a matrix of M x N elements (M rows, N columns), return all elements of the matrix in diagonal order as shown in the below image.

Example:

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

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& matrix) {
        if(matrix.empty())
            return vector<int>();
        int rowNum=matrix.size();
        int colNum=matrix[0].size();
        int i=0,j=0;
        vector<int> res;
        int num=(matrix.size())*matrix[0].size();
        int count=0,d=1;
        while(count<num)
        {
            res.push_back(matrix[i][j]);
            count+=1;
            i-=d;
            j+=d;
            if(i==rowNum)
            {
                i=rowNum-1;
                j+=2;
                d=-d;
            }
            if(j==colNum)
            {
                j=colNum-1;
                i+=2;
                d=-d;
            }
            if(i<0)
            {
                i=0;
                d=-d;
            }
            if(j<0)
            {
                j=0;
                d=-d;
            }
        }
        return res;
    }
};