facebook/meta 2D spiral array

2D_spiral_array.py

"""
Find the pattern and complete the function:
int[][] spiral(int n);
where n is the size of the 2D array.

input = 3
123
894
765

input = 4
01 02 03 04
12 13 14 05
11 16 15 06
10 09 08 07
"""


import numpy as np

def isInvalid(n, row, col, matrix):    
    if row < 0 or col < 0:
        return True
    if row == n or col == n:
        return True
    if matrix[row, col] != 0:
        return True

def turn_right(turns):
    direction = turns % 4
    
    if direction == 1: # down
        row_step = 1
        col_step = 0
    if direction == 2: # left
        row_step = 0
        col_step = -1
    if direction == 3: # up
        row_step = -1
        col_step = 0
    if direction == 0: # right
        row_step = 0
        col_step = 1
        
    return row_step, col_step
    
def step_foward(row, col, r_step, c_step):
    return row+r_step, col+c_step
    
def spiral(n):
    matrix = np.zeros((n,n))
    row, col, curr = 0, 0, 1
    r_step, c_step = 0, 1
    turns = 1
    limit = n*n

    # 1: row constant, col+=1
    # 2: col constant, row+=1
    # 3: row constant, col-=1
    # 4: col constatn, row-=1

    while curr <= limit: 

        # print(f'row {row}, col {col}, curr {curr}')
        matrix[row, col] = curr
        curr += 1
    
        if isInvalid(n, row+r_step, col+c_step, matrix):
            r_step, c_step = turn_right(turns)
            # print(f'r_step {r_step}, c_step {c_step}')
            turns += 1    

        row, col = step_foward(row, col, r_step, c_step)

    return matrix
    
print(spiral(9))