코드

처음에 작성한 코드 => 시간 초과

from collections import deque
import sys
input = sys.stdin.readline

n, m, p = map(int, input().split())
s = [0]+list(map(int, input().split()))
board = [list(input()) for _ in range(n)]
ans = [0]*(p+1)

dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

castle = [[] for _ in range(p+1)]
# 처음 성의 위치
for i in range(n):
    for j in range(m):
        if board[i][j] != '.' and board[i][j] != '#':
            castle[int(board[i][j])].append((i, j))
            ans[int(board[i][j])] += 1


def expand(player):
    q = deque()
    visited = [[int(1e9)]*m for _ in range(n)]
    for x, y in castle[i]:
        q.append((x, y))
        visited[x][y] = 0
    castle[i].clear()
    while q:
        x, y = q.popleft()
        for d in range(4):
            nx = x+dx[d]
            ny = y+dy[d]
            if 0 <= nx < n and 0 <= ny < m and board[nx][ny] == '.':
                if visited[nx][ny] > visited[x][y]+1:
                    board[nx][ny] = str(player)
                    castle[player].append((nx, ny))
                    ans[int(board[nx][ny])] += 1
                    visited[nx][ny] = visited[x][y]+1
                    if visited[x][y]+1 < s[player]:
                        q.append((nx, ny))


while True:
    for i in range(1, p+1):
        if len(castle[i])>0:
            expand(i)
    is_end = True
    for i in range(1, p+1):
        if len(castle[i]) > 0:
            is_end = False
            break
    if is_end:
        break

print(*ans[1:])

castle[i]는 각 턴에 확장된 칸들의 정보가 계속 담기고 이 정보를 바탕으로 계속하여 BFS를 진행하는 코드인데, 80%에서 계속 시간 초과가 발생해서 결국 구글링해서 찾아낸 다른 방법으로 바꿨다.

from collections import deque
import sys
input = sys.stdin.readline

n, m, p = map(int, input().split())
s = [0]+list(map(int, input().split()))
board = [list(input()) for _ in range(n)]
ans = [0]*(p+1)

castle = [deque() for _ in range(p+1)]
# 처음 성의 위치
for i in range(n):
    for j in range(m):
        if board[i][j] != '.' and board[i][j] != '#':
            castle[int(board[i][j])].append((i, j))
            ans[int(board[i][j])] += 1

# BFS
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

while True:
    is_end = True
    for player in range(1, p+1):
        if not castle[player]:
            continue
        q = castle[player]
        for _ in range(s[player]):  # 한 턴에 이동할 수 있는 거리
            if not q:   # 더이상 이동 가능한 칸이 없는 경우
                break
            for _ in range(len(q)):
                x, y = q.popleft()
                for d in range(4):
                    nx = x+dx[d]
                    ny = y+dy[d]
                    if 0 <= nx < n and 0 <= ny < m and board[nx][ny] == '.':
                        board[nx][ny] = str(player)
                        q.append((nx, ny))
                        ans[player] += 1
                        is_end = False
    if is_end:
        break


print(*ans[1:])

첫번째 코드와는 다르게 턴 한번에는 s[player]의 값만큼의 거리까지만 BFS 탐색이 가능하도록 했다. 이렇게 하면 BFS를 하기 위해 불필요하게 덱을 따로 만든다거나 이미 지나간 칸을 체크하기 위해 visited를 사용할 필요가 없기에 시간을 단축시킬 수 있는 것 같다.

저작자표시 (새창열림)