Diagonal Traverse

class Solution:
    def findDiagonalOrder(self, matrix):
        """
        :type matrix: List[List[int]]
        :rtype: List[int]
        """
        # 00 (right: col+1) //0 
        # 01, 10 (left: row+1, col-1) -> (if col==0: row+1) //1
        # 20, 11, 02 (right: row-1, col+1) -> if(row==0: col+1) //2
        # 03, 12, 21, 30 (left_middle: row+1, col-1) -> if col==0: col+1 //3--> col+1, (right_middle,,, row+1) ::> 
        
        # -----
        # 31, 22, 13 (right_middle: row-1, col+1) -> if row==1, row+1(state_count - len(matrix))
        # 23, 32, if (left_middle: row+1, col-1) if col==2: col+1..
        # 33 ...if (iterate_nums == len(matrix) ** 2) <- exponential.
        output = []
        row, col = 0, 0
        state = 0
        state_count = 1
        for index in range(len(matrix)**2):
            print(row)
            print(col)
            print('--')
            output.append(matrix[row][col])
            # rightstate
            if state == 0:
                if row == 0:
                    state = 1
                    # track the state change
                    state_count = state_count + 1
                    if state_count > len(matrix):
                        state = 3
                        # row = row +1
                        col = col - 1
                        row = row +1
                    row = row + 1 # preventing from increasing row
                col = col + 1
                row = row - 1
            # left state
            elif state == 1:
                if col == 0:
                    state = 0
                    # track the state change
                    state_count = state_count + 1
                    if state_count > len(matrix):
                        state = 2
                        # col = col +1
                        row = row - 1
                        col = col +1
                    col = col + 1 # preventing from increasing row
                row = row + 1
                col = col - 1
            # right middle state
            elif state == 2:
                # 31, 22, 13 (right_middle: row-1, col+1) -> if row==1, row+1(state_count - len(matrix))
                if row == state_count - len(matrix):
                    state_count = state_count + 1
                    state = 3
                    # incremenet row
                    row = row + 2
                    col = col - 1
                row = row - 1
                col = col + 1
                
            # left middle state
            elif state == 3:
                if col == state_count - len(matrix):
                    state_count = state_count + 1
                    state = 2
                    # col = col + state_count - len(matrix)
                    # row = row - 1
                    # incremenet col
                    col = col + 2
                    row = row - 1
                row = row + 1
                col = col - 1
        return output