A single pass solution using BFS-like approach. Algorithm-
- Start by placing coordinates
(0,0)in a queue for BFS. - In each iteration, for the
frontelement in queue push the index for cell just below if the current cell is the first in its row. - Then push the index for right neighbour cell (if exists)
class Solution {
public:
vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
int m = nums.size();
vector<int> ans;
queue<pair<int, int>> q;
q.push({0,0}); // first row, first cell
// BFS
while (!q.empty()) {
pair<int, int> p = q.front();
q.pop();
ans.push_back(nums[p.first][p.second]);
// insert the element below, if in first column
if (p.second == 0 && p.first+1 < m) q.push({p.first+1, p.second});
// insert the right neighbour, if exists
if (p.second+1 < nums[p.first].size())
q.push({p.first, p.second+1});
}
return ans;
}
};