<p>Given a 2D integer array <code>nums</code>, return <em>all elements of </em><code>nums</code><em> in diagonal order as shown in the below images</em>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2020/04/08/sample_1_1784.png" style="width: 158px; height: 143px;" />
<pre>
<strong>Input:</strong> nums = [[1,2,3],[4,5,6],[7,8,9]]
<strong>Output:</strong> [1,4,2,7,5,3,8,6,9]
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2020/04/08/sample_2_1784.png" style="width: 230px; height: 177px;" />
<pre>
<strong>Input:</strong> nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
<strong>Output:</strong> [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= nums.length &lt;= 10<sup>5</sup></code></li>
	<li><code>1 &lt;= nums[i].length &lt;= 10<sup>5</sup></code></li>
	<li><code>1 &lt;= sum(nums[i].length) &lt;= 10<sup>5</sup></code></li>
	<li><code>1 &lt;= nums[i][j] &lt;= 10<sup>5</sup></code></li>
</ul>

class Solution {
public:
    vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
        vector<int> ans;
        priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<tuple<int, int, int>>> q;
        q.push({0, 0, 0});
        int n = nums.size();
        vector<int> m(n);
        for(int i = 0; i < n; i++){
            m[i] = nums[i].size();
        }
        vector<int> dr = {1, 0};
        vector<int> dc = {0, 1};
        vector<vector<bool>> vis(n); 
        for(int i = 0; i < n; i++){
            vis[i].resize(m[i], false); 
        }
        vis[0][0] = true;
        vector<int> diag;

        while(!q.empty()){
            int l = q.size();
            diag.clear();
            while(l--){
                tuple<int, int, int> t = q.top();
                int r = get<1>(t);
                int c = get<2>(t);
                int sum = get<0>(t);
                diag.push_back(nums[r][c]);
                q.pop();
                if(r == n-1 && c == m[r]){
                    return ans;
                }
                for(int i = 0; i < 2; i++){
                    int row = r +dr[i];
                    int col = c + dc[i];
                    if(row>=0 && row < n && col >=0 && col < m[row] && vis[row][col] == false){
                        q.push({row+col, row, col});
                        vis[row][col] = true;
                    }
                }
            }
            reverse(diag.begin(), diag.end());
            ans.insert(ans.end(), diag.begin(), diag.end());            
        }
        return ans;
    }
};
80238