Ruudukko

def generate_empty_grid(n):
    grid = []
    for row in range(0, n):
        grid.insert(row, [0] * n)
    return grid
 
def get_int(n_col, n_offset, n):
    v = n_col + n_offset + 1
    if v > n:
        return v - n
    else:
        return v
 
def get_new_offset(cur_offset, n):
    new_offset = cur_offset + 1
    if new_offset > n:
        new_offset = 0
    return new_offset
 
def main(n):
    grid = generate_empty_grid(n)
 
    n_offset = 0
    for n_row, row in enumerate(grid):
        for n_col in range(0, len(row)):
            v = get_int(n_col, n_offset, n)
            row[n_col] = v
            if n_col == len(row) - 1:
                print(v, end = '')
            else:
                print(v, end = ' ')
        n_offset = get_new_offset(n_offset, n)
        print("\n", end = '')
 
main(int(input()))

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: WRONG ANSWER

Test 4

Verdict: WRONG ANSWER

Test 5

Verdict: WRONG ANSWER

Test 6

Verdict: WRONG ANSWER