1424. Diagonal Traverse II

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) 
    {
        int n = nums.size();
        int count = 0;

        ArrayList<List<Integer>> list = new ArrayList<>();

        for (int i = 0; i < n; i++) 
        {
            List<Integer> row = nums.get(i);

            for (int j = 0; j < row.size(); j++) 
            {
                int idx = i + j;

                if (list.size() < idx + 1) 
                {
                    list.add(new ArrayList<>());
                }

                list.get(idx).add(row.get(j));
                
                count ++;
            }
        }

        int[] res = new int[count];
        int idx = 0;

        for (List<Integer> x : list)
        {
            for (int i = x.size() - 1; i >= 0; i--) 
            {
                res[idx++] = x.get(i); 
            }

        }
        
        return res;

        
    }
}