AcWing 3208 zigzag scan

Title description:

In the image coding algorithm, a given square matrix needs to be Zigzag Scan.

Given an n×n matrix, the zigzag scanning process is shown in the figure below:

For the following 4×4 matrix,

1 5 3 9
3 7 5 6
9 4 6 4
7 3 1 3

Z-shaped after it was scanned in a sequence of length 16: 1 5 3 9 7 3 9 5 4 7 3 6 6 4 1 3.

Please implement a zigzag scanning program, given an n×n matrix, output the result of zigzag scanning on this matrix.

Input format

The first line of input contains an integer n, which represents the size of the matrix.

Each line from the second line to the n+1th line of the input contains n positive integers, separated by spaces, representing the given matrix.

Output format

Output one line, containing n×n integers, separated by spaces, representing the result of the input matrix after zigzag scanning.

data range

1≤n≤500,
matrix elements are positive integers not exceeding 1000.

Input sample:

4
1 5 3 9
3 7 5 6
9 4 6 4
7 3 1 3

Sample output:

1 5 3 9 7 3 9 5 4 7 3 6 6 4 1 3

#include <iostream>
#include <cstdio>

using namespace std;
const int MAX = 509;

int a[MAX][MAX];
int ans[MAX * MAX];

int n;

bool check(int x, int y)
{
    if(x >= 0 && x < n && y >= 0 && y < n)
        return true;
    return false;
}

int main()
{
    scanf("%d", &n);

    for(int i = 0; i < n; i++)
    {
        for(int j = 0; j < n; j++)
            scanf("%d", &a[i][j]);
    }

    int N = n * n;
    int cnt = 1, k = 0;
    bool down = false; // down为false向右走，否则向下走
    int row = 0, col = 0;
    bool rightup = true; // true向右上方走，false向左下方走
    bool f = false;

    while(k < N)
    {
        for(int i = 0; i < cnt; i++)
        {
            ans[k++] = a[row][col];
            // printf("(%d, %d)  a = %d\n", row, col, a[row][col]);
            if(check(row - 1, col + 1) && rightup)
            {
                row--;
                col++;
            }
            if(check(row + 1, col - 1) && !rightup)
            {
                row++;
                col--;
            }
        }
        // printf("\n");

        if(cnt >= n)
            f = true;

        if(!f)
            cnt++;
        else
            cnt--;

        if(!down)
        {
            if(check(row, col + 1)) // 可以向右走
            {
                col ++;
                down = true;
            }
            else
                row++;
        }

        else
        {
            if(check(row + 1, col)) // 可以向下走
            {
                row++;
                down = false;
            }
            else
                col++;
        }
        rightup = !rightup;
    }

    N = n * n;
    for(int i = 0; i < N; i++)
        printf("%d ", ans[i]);

    return 0;
}