Diagonal Traverse II

class Solution {
    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        int count = 0;
        for(int i = 0; i<nums.size(); i++){
            count+=nums.get(i).size();
        }
        int[] res = new int[count];
        int[][] u = new int[count][3];
        count = 0;
        for(int i = nums.size()-1; i>=0; i--){
            for(int j = 0; j<nums.get(i).size(); j++){
                u[count++] = new int[]{i+j, j, nums.get(i).get(j)};
            }
        }
        Arrays.sort(u, new Comparator<int[]>() {
				public int compare(int[] a, int[] b) {
					if(a[0] != b[0])return a[0] - b[0];
					return a[1] - b[1];
				}
			});
        //System.out.println(Arrays.deepToString(u));
        for(int i = 0; i<res.length; i++){
            res[i] = u[i][2];
        }
        return res;
    }
}