[ BOJ / Python ] 16957번 체스판 위의 공

황승환·2022년 7월 24일
0

Python

목록 보기
390/498


이번 문제는 BFS와 DP를 통해 해결하였다. 처음 풀이에서는 오로지 BFS만 사용하여 각 격자에 존재하는 모든 구슬에 대하여 BFS탐색을 실시하였다. 당연히 답은 잘 도출되었지만, 70%에서 시간초과가 발생하였다.

이를 줄일 수 있는 방법은 메모라이제이션이라 생각하였고, 모든 격자의 경로가 정해져있음을 깨달았다. 그렇기 때문에 각 격자의 최종 도착지를 저장하는 리스트를 만들고, 최종 도착지를 찾는 함수를 통해 최종 도착지 리스트를 갱신하였다. 이 과정에서 중간에 들리는 좌표의 최종 도착지가 정의되어 있을 경우, 같은 좌표가 최종 도착지가 되므로 바로 최종 도착지로 갱신해주었다.

Code

처음 제출 코드 (BFS) - 시간초과

from collections import deque
r, c = map(int, input().split())
grid = [list(map(int, input().split())) for _ in range(r)]
marbles = [[1 for _ in range(c)] for _ in range(r)]
dy, dx = [-1, -1, 0, 1, 1, 1, 0, -1], [0, 1, 1, 1, 0, -1, -1, -1]
def bfs(sy, sx):
    q = deque()
    q.append((sy, sx))
    while q:
        y, x = q.popleft()
        mn, idx = 1e9, []
        for i in range(8):
            ny, nx = y+dy[i], x+dx[i]
            if 0 <= ny < r and 0 <= nx < c:
                if grid[ny][nx] <= grid[y][x] and mn >= grid[ny][nx]:
                    mn = grid[ny][nx]
                    idx = [ny, nx]
        if mn == 1e9 and not idx:
            return [y, x]
        else:
            q.append((idx[0], idx[1]))
new_marbles = [[0 for _ in range(c)] for _ in range(r)]
for i in range(r):
    for j in range(c):
        if marbles[i][j]:
            nxt_idx = bfs(i, j)
            new_marbles[nxt_idx[0]][nxt_idx[1]] += marbles[i][j]
for i in range(r):
    print(*new_marbles[i])

정답 코드 (BFS + DP)

from collections import deque
r, c = map(int, input().split())
grid = [list(map(int, input().split())) for _ in range(r)]
marbles = [[1 for _ in range(c)] for _ in range(r)]
parent = [[[] for _ in range(c)] for _ in range(r)]
dy, dx = [-1, -1, 0, 1, 1, 1, 0, -1], [0, 1, 1, 1, 0, -1, -1, -1]
def bfs(sy, sx):
    q = deque()
    q.append((sy, sx))
    while q:
        y, x = q.popleft()
        mn, idx = 1e9, []
        for i in range(8):
            ny, nx = y+dy[i], x+dx[i]
            if 0 <= ny < r and 0 <= nx < c:
                if grid[ny][nx] <= grid[y][x] and mn >= grid[ny][nx]:
                    mn = grid[ny][nx]
                    idx = [ny, nx]
        if mn == 1e9 and not idx:
            return [y, x]
        if parent[idx[0]][idx[1]]:
            return parent[idx[0]][idx[1]]
        else:
            q.append((idx[0], idx[1]))
for i in range(r):
    for j in range(c):
        parent[i][j] = bfs(i, j)
new_marbles = [[0 for _ in range(c)] for _ in range(r)]
for i in range(r):
    for j in range(c):
        if marbles[i][j]:
            nxt_idx = parent[i][j]
            new_marbles[nxt_idx[0]][nxt_idx[1]] += marbles[i][j]
for i in range(r):
    print(*new_marbles[i])

profile
꾸준함을 꿈꾸는 SW 전공 학부생의 개발 일기

0개의 댓글