I tried to solve AtCoder's depth-first search (DFS) in Python (result: TLE ...)

class Dfs():
    def dfs(self, h, w, maze, init):
        search_stack = [init] #Store initial coordinates in search stack
        fin = [] #Define the completion stack
        while search_stack:
            latest = search_stack.pop() #Extract your location from the search stack
            if maze[latest[0]][latest[1]] == "g": #Output Yes if your current location is Goal
                print("Yes")
                exit()
            else:
                fin.append(latest) #Store the searched coordinates on the completion stack
                candi = [(latest[0]-1, latest[1]), (latest[0]+1, latest[1]), (latest[0], latest[1]-1), (latest[0], latest[1]+1)] #Define coordinate candidates
                for (ny, nx) in candi:
                    if self.check_range(h, w, ny, nx): #Determine if the coordinate candidates are out of range
                        if self.check_wall(maze, ny, nx): #Determine if the coordinate candidate is a wall
                            if not (ny, nx)  in search_stack and not (ny, nx) in fin: #Determine if coordinate candidates are included in the search stack and completion stack
                                search_stack.append((ny, nx)) #Store in search stack
                                continue
        print("No")                    
        
    def check_range(self, h, w, y, x):
        if x > w-1 or x < 0:
            return False
        elif y > h-1 or y < 0:
            return False
        else:
            return True
    
    def check_wall(self, maze, y, x):
        if maze[y][x] == "#":
            return False
        else: 
            return True

if __name__ == "__main__":
    h,w = map(int,input().split())
    maze = [list(input()) for i in range(h)]
    for y, row in enumerate(maze):
        try:
            init = (y, row.index("s")) #Search for a starting point
            d = Dfs()
            d.dfs(h, w, maze, init)
            break
        except ValueError:
            pass

Recommended Posts