O(n) intuitive Solution c++

[KayMK11]

vector<vector<int> > Solution::diagonal(vector<vector<int> > &A) {
int n = A.size() - 1;
int i = 0, j = 0, k = 0;
vector<vector<int>> ans(2 * n + 1);
while (k < 2 * n + 1)
{
    if ((i >= 0 && i <= n) && (j >= 0 && j <= n))
    {
        ans[k].push_back(A[i][j]);
        i++;
        j--;
    }
    else
    {
        k++;
        if (k < n)
        {
            i = 0;
            j = k;
        }
        else
        {
            i = k - n;
            j = n;
        }
    }
}
return ans;

}