[Python] BOJ 백준 21922 학부연구생 민상

문제

학부 연구생으로 새로 연구실에 들어온 민상이는 사용할 자리를 정하려고 한다.

연구실은 격자 모양으로 되어있고 에어컨에서 바람이 상,하,좌,우 4방향으로 분다. 물론 에어컨이 위치한 곳에도 바람이 분다.

민상이는 더위를 많이 타서 에어컨 바람이 지나가는 곳 중 하나를 선택하여 앉으려고 한다.

연구실에는 다양한 물건들이 있어 바람의 방향을 바꾼다.

연구실에 있는 물건의 종류는 총 4가지가 있다. 아래 화살표의 의미는 바람이 각 물건에서 바람의 이동을 표시한 것이다.

연구실 어디든 민상이가 앉을 수 있는 자리이다. 즉 에어컨이 위치한 자리와 물건이 있는 자리 모두 앉을 수 있다.

민상이가 원하는 자리는 몇 개 있는지 계산해주자.

풀이

from collections import deque
import sys

si = sys.stdin.readline

n, m = map(int, si().split())
graph = [list(map(int, si().split())) for _ in range(n)]
visited = [[0] * m for _ in range(n)]


queue = deque()

for i in range(n):
    for j in range(m):
        if graph[i][j] == 9:
            queue.append((i, j))


def bfs(g, q):
    xd = [-1, 1, 0, 0]
    yd = [0, 0, -1, 1]
    while q:
        x, y = q.popleft()
        for i in range(4):
            r, c = x, y
            nx, ny = xd[i], yd[i]
            while 0 <= r < n and 0 <= c < m:
                visited[r][c] = 1
                if (graph[r][c] == 1 and nx == 0) or (graph[r][c] == 2 and ny == 0):
                    break
                elif graph[r][c] == 3:
                    nx, ny = -ny, -nx
                elif graph[r][c] == 4:
                    nx, ny = ny, nx
                r += nx
                c += ny
    return sum(sum(visited, []))


print(bfs(graph, queue))

이 문제는  너비 탐색을 이용해서 풀었다. 

진행하던 방향에 따라서 어떻게 달라지는 지에 대해서 공식을 세워두고, 분기를 나누어서 처리했다. 계속해서 시간 초과가 났었는데,

알고보니 sum(sum(visited, [])) 부분에서 시간 초과가 났다. 단순히 count로만 바꿔도 python으로 풀 수 있다. 

from collections import deque
import sys

si = sys.stdin.readline

n, m = map(int, si().split())
visited = [[0] * m for _ in range(n)]


queue = deque()

graph = []
for i in range(n):
    line = list(map(int, si().split()))
    for j in range(m):
        if line[j] == 9:
            queue.append((i, j))
            visited[i][j] = 1
    graph.append(line)


def bfs(g, q):
    xd = [-1, 1, 0, 0]
    yd = [0, 0, -1, 1]
    while q:
        x, y = q.popleft()
        for idx in range(4):
            nx, ny = xd[idx], yd[idx]
            r, c = x + nx, y + ny
            while 0 <= r < n and 0 <= c < m:
                visited[r][c] = 1
                if g[r][c] == 9:
                    break
                if g[r][c] == 3:
                    nx, ny = -ny, -nx
                elif g[r][c] == 4:
                    nx, ny = ny, nx
                elif (g[r][c] == 1 and nx == 0) or (g[r][c] == 2 and ny == 0):
                    break
                r += nx
                c += ny
    answer = 0
    for ans in visited:
        answer += ans.count(1)
    return answer


print(bfs(graph, queue))