之字形打印矩阵

链接
给定一个矩阵matrix，按照“之”字形的方式打印这个矩阵，如样例所示。

import java.util.*;

public class Main {

    private static int[] solve(int[][] mat) {
        int n = mat.length;
        int m = mat[0].length;
        int[] ret = new int[n * m];
        int index = 0;

        int row1 = 0, col1 = 0, row2 = 0, col2 = 0;
        boolean dir = false;

        while (row1 < n && col2 < m) {
            if (dir) {
                int row = row1, col = col1;
                while (row <= row2 && col >= col2) {
                    ret[index++] = mat[row++][col--];
                }
            } else {
                int row = row2, col = col2;
                while (row >= row1 && col <= col1) {
                    ret[index++] = mat[row--][col++];
                }
            }
            dir = !dir;
            if (col1 == m - 1) {
                row1++;
            } else {
                col1++;
            }
            if (row2 == n - 1) {
                col2++;
            } else {
                row2++;
            }
        }

        return ret;
    }

    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        while (in.hasNext()) {
            int n = in.nextInt();
            int m = in.nextInt();
            int[][] mat = new int[n][m];
            for (int i = 0; i < n; ++i) {
                for (int j = 0; j < m; ++j) {
                    mat[i][j] = in.nextInt();
                }
            }
            int[] ret = solve(mat);
            System.out.print(ret[0]);
            for (int i = 1; i < ret.length; ++i) {
                System.out.print(" " + ret[i]);
            }
            System.out.println();
        }
    }
}