<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>

# include <bits/stdc++.h>

class Solution {
public:

    void printAns(vector<int> &ans) {
        for (auto num : ans) {
            cout << num << " ";
        }
        cout << endl;
    }

    void printRows(set<int> &rows) {
        for (auto row : views::reverse(rows)) {
            cout << row << " ";
        }
        cout << endl;
    }

    void addDiagonal(vector<int> &ans, set<int> &rows, vector<vector<int>>& nums, int startRow, int startCol) {
        rows.insert(startRow);

        // cout << "start row: " << startRow << " start col: " << startCol << endl;
        // printRows(rows);
        // printAns(ans);
        vector<int> rowsToErase;

        for (auto row : views::reverse(rows)) {
            int col = startRow - row + startCol;
            if (col >= nums[row].size()) {
                rowsToErase.emplace_back(row);
                continue;
            }
            ans.emplace_back(nums[row][col]);
            if (startRow == 2 and startCol == 1) {
                printAns(ans);
            }
        }

        for (int row : rowsToErase) {
            rows.erase(row);
        }
    }
    vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
        // we essentially do the same thing
        // we want to iterate over each diagonal in order
        // we start in column 0, the row we start at moves down one each time, in each previous row, we visit column startRow - Row
        // column we visit for each row increases each time
        // therefore the second we try to visit a column in a row, and it is not valid, we never have to visit that row again, and before this moment all visits to that row will be valid

        // we store a bst of valid row indices
        // each time we iterate over the still valid row indices in order from greatest to least, check the corresponding column, if existent, add to answer, otherwise, remove that row from bst
        // with each diagonal we add new row to bst
        // m(n + logm) time where m is the number of lists and m is the maximum length of a list
        // but more precisely it take O(mlogm + number of total elements) time

        set<int> rows;
        vector<int> ans;
        int maxLen = 0;

        for (int i = 0; i < nums.size(); i++) {
            maxLen = max(maxLen, int(nums[i].size()));

            addDiagonal(ans, rows, nums, i, 0);
        }

        // cout << "here" << endl;
        // cout << maxLen << endl;
        // cout << nums.size() - 1 << endl;

        for (int j = 1; j < maxLen; j++) {
            // cout << "here" << endl;
            addDiagonal(ans, rows, nums, nums.size() - 1, j);
        }
        return ans;
    }
};
95091