1424. Diagonal Traverse II 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
    private static class Cell implements Comparable<Cell> {
        private final int rowIndex;
        private final int colIndex;
        private final int value;

        public Cell(int rowIndex, int colIndex, int value) {
            this.rowIndex = rowIndex;
            this.colIndex = colIndex;
            this.value = value;
        }

        @Override
        public int compareTo(Cell o) {
            int sum1 = rowIndex + colIndex;
            int sum2 = o.rowIndex + o.colIndex;
            int diff = Integer.compare(sum1, sum2);
            if (diff != 0) {
                return diff;
            } else {
                return Integer.compare(o.rowIndex, rowIndex);
            }
        }
    }

    public int[] findDiagonalOrder(List<List<Integer>> nums) {
        List<Cell> cellList = new ArrayList<>();
        for (int i = 0; i < nums.size(); i++) {
            List<Integer> row = nums.get(i);
            for (int j = 0; j < row.size(); j++) {
                cellList.add(new Cell(i, j, row.get(j)));
            }
        }

        Collections.sort(cellList);

        int[] res = new int[cellList.size()];
        for (int i = 0; i < cellList.size(); i++) {
            res[i] = cellList.get(i).value;
        }
        return res;
    }
}

References

1424. Diagonal Traverse II