digonal sort

class Solution {
    public int[][] diagonalSort(int[][] mat) {
        int m = mat.length;
        int n = mat[0].length;

        // Diagonals starting from first row
        for (int col = 0; col < n; col++) {
            sortDiagonal(mat, 0, col, m, n);
        }

        // Diagonals starting from first column (excluding the top-left one again)
        for (int row = 1; row < m; row++) {
            sortDiagonal(mat, row, 0, m, n);
        }

        return mat;
    }

    private void sortDiagonal(int[][] mat, int row, int col, int m, int n) {
        List<Integer> diag = new ArrayList<>();
        int r = row, c = col;

        // Collect elements along the diagonal (down-right)
        while (r < m && c < n) {
            diag.add(mat[r][c]);
            r++;
            c++;
        }

        // Sort the diagonal
        Collections.sort(diag);

        // Put sorted values back
        r = row;
        c = col;
        int idx = 0;
        while (r < m && c < n) {
            mat[r][c] = diag.get(idx++);
            r++;
            c++;
        }
    }
}