<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:
    struct Entry
    {
        int x;
        int y;
        int val;
        long long position;
    
        void computePosition()
        {
            position = static_cast<long long>(x + y + 1) * (x + y + 2) / 2 - y;
        }
        bool operator<(Entry const& arg) const
        {
            return position < arg.position;
        }
    };
    vector<int> findDiagonalOrder(vector<vector<int>>& nums) {
        //For an infinite square grid with an upper left corner,
        //we have a function that takes 2D position to ordinal number:
        //(0,0) -> 1
        //(0,1) -> 2
        //(1,0) -> 3
        //(0,2) -> 4
        //(1,1) -> 5
        //(2,0) -> 6
        //etc.
        //x+y = 0 -> [1]
        //x+y = 1 -> [2,3]
        //x+y = 2 -> [4,6]
        //The top of the range are the triangular numbers
        //T = (x + y + 1)(x + y + 2) / 2
        //Therefore f(x,y) = (x + y + 1)(x + y + 2) / 2 - y

        vector<Entry> entries;
        for (int y = 0; y < nums.size(); ++y)
        {
            for (int x = 0; x < nums[y].size(); ++x)
            {
                Entry e;
                e.x = x;
                e.y = y;
                e.val = nums[y][x];
                e.computePosition();
                entries.push_back(e);
            }
        }

        sort(entries.begin(), entries.end());
        
        vector<int> result;
        result.reserve(entries.size());
        for (auto const& e : entries)
        {
            result.push_back(e.val);
        }
        return result;
    }
};
87228