Zigzag Conversion

/*
 * Problem : Zigzag text.
 * Input   : A text and a row count.
 * Output  : The text read row by row.
 * Approach: Bounce between rows.
 * Time    : O(n)
 * Space   : O(n)
 * Java    : 8 or newer.
 */
import java.util.Scanner;

public class Main {
    // This method solves the task. It receives the input and returns the answer.
    static String solve(String a, String b) {
        // Step 1: With one row, or with more rows than letters, the text does not change.
        int rows = Integer.parseInt(b);
        if (rows == 1 || rows >= a.length()) {
            return a;
        }
        // Step 2: One StringBuilder per row collects the letters of that row.
        StringBuilder[] lines = new StringBuilder[rows];
        for (int i = 0; i < rows; i++) {
            lines[i] = new StringBuilder();
        }
        // Step 3: Walk through the text. Turn around at the first and the last row.
        int row = 0;
        int step = 1;
        for (int i = 0; i < a.length(); i++) {
            lines[row].append(a.charAt(i));
            if (row == 0) {
                step = 1;
            } else if (row == rows - 1) {
                step = -1;
            }
            row += step;
        }
        // Step 4: Join the rows from top to bottom.
        StringBuilder result = new StringBuilder();
        for (StringBuilder line : lines) {
            result.append(line);
        }
        return result.toString();
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        String a = scanner.next();
        String b = scanner.next();
        System.out.println(solve(a, b));
    }
}