<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:
    # [1,2,3,10,   e4],
    # [4,5,6, 100, e3],
    # [7,8,9, 200, e2]
    # [a,b,c, d, e]

    # ind_row = 2 
    # ind_col = 0
    # maxC = 3

    """
    traversal = [
        0: [1, ]
        1: [2, 4]
        2: [3, 5, 7]

        3: 
        4: 

    ]
    """

    def findDiagonalOrder(self, nums: List[List[int]]) -> List[int]:

        R = len(nums)
        maxC = max(len(_) for _ in nums)

        traversal = {}
        for i in range(R + maxC):
            traversal[i] = []
        
        for ind_row, row in enumerate(nums):
            for ind_col, cellValue in enumerate(row):
                # part of the first pass 
                if ind_col <= maxC - ind_row - 1:
                    traversal[ind_row + ind_col].append(cellValue)
                else:
                    # how to determine the traversal index? 
                    # we need to take into account ind_row as well 
                    # ind_col - (R - ind_row - 1)
                    # r1, c1 should have the same index as r1 + 1, c1 - 1
                    # f(r1, c1) = f(r1 + 1, c1 -1) = f(r1 + 2, c1 - 2) = f(r1 + 3, c1 - 3)
                    # f(r, c) = c - (R - r - 1)
                    # f(3, 0) = 0 - (4 - 3 - 1) = 0
                    # f(2, 1) = 1 - (4 - 2 - 1) = 0 
                    # f(1, 2) = 2 - (4 - 1 - 1) = 0

                    # f(3, 1) = 1 - (4 - 3 - 1) = 1
                    # f(2, 2) = 2 - (4 - 2 - 1) = 1
                    # f(1, 3) = 3 - (4 - 1 -1) = 1 
                    trav_ind = ind_col - (R - ind_row - 1)
                    traversal[R + trav_ind].append(cellValue)

        result = []

        for trav_ind in range(R + maxC):
            result.extend(traversal[trav_ind][::-1])
        
        return result

        


        
46631