[백준] 16236 아기상어 (Python/파이썬)

<문제 링크>

https://www.acmicpc.net/problem/16236

16236번: 아기 상어

 

<풀이 전략>

1. 구현 문제이다. 나같은 경우는 처음에 최대 O(N^6)의 코드로 pypy(1076ms)로 통과 (python 시간초과)하고 이후에 jh05013님의 풀이를 참고하여 훨씬 빠른 풀이(python: 288ms)로 다시 풀었는데 두 가지 모두 소개해보겠다.
2. 가장 중요한 것은 상어가 다음 먹이를 찾아가는 과정을 구현하는 것인데, 여기서 '가장 거리가 가까우며,  거리가 같을 경우 왼쪽 상단에 있는 먹이'를 고르는 조건을 구현하는 것이 까다롭다.  
   1. 나같은 경우는 주어진 배열을 (0, 0)부터 (N-1, N-1)까지 '해당 좌표의 먹이 -> 상어'로 bfs 탐색을 해서 거리가 더 가까울 때만 다음 먹이 정보를 갱신하는 방식을 사용했고,
   2. 빠른 풀이같은 경우는 '상어 -> 모든 먹이' 로 한번의 bfs탐색을 해서 모든 먹이의 거리를 구한 다음에, 해당 먹이의 거리와 좌표를 별도의 배열에 담아 거리를 기준으로 min 값으로 다음 먹이의 정보를 구했다. 이렇게 할 경우 거리 -> 행 -> 열 순으로 구한 최솟값이 나오기 때문에 한번의 bfs탐색으로 끝낼 수 있다.
3. 그 외 자잘한 구현은 특별히 어려운 게 없어서 아래 코드 주석을 참고하면 된다.
4. 참고로 해당 문제는 삼성 기출인데 내가 알기로 삼성 시험은 pypy를 사용하기 때문에 조금 무식하게 첫번째 풀이로 해도 통과는 문제 없다. 

<정답코드>

* 빠른 풀이

import sys, os, io, atexit

input = lambda: sys.stdin.readline().rstrip('\r\n')
stdout = io.BytesIO()
sys.stdout.write = lambda s: stdout.write(s.encode("ascii"))
atexit.register(lambda: os.write(1, stdout.getvalue()))

from collections import deque

# 델타 탐색 방향 설정
dr = [-1, 1, 0, 0]
dc = [0, 0, -1, 1]

# 아기 상어 먹이 거리 탐색 bfs
def bfs(shark_r, shark_c):
    # visited: 방문여부 & 먹이까지의 거리 표시할 배열
    visited = [[inf for _ in range(N)] for _ in range(N)]
    visited[shark_r][shark_c] = 0

    queue = deque()
    queue.append((shark_r, shark_c))

    while queue:
        r, c = queue.popleft()
        for i in range(4):
            nr = r + dr[i]
            nc = c + dc[i]
            # 조건 맞으면 거리 표시하면서 탐색
            if 0 <= nr < N and 0 <= nc < N and visited[nr][nc] == inf:
                # 먹이가 상어의 현재 크기보다 작거나 같은 경우
                if arr[nr][nc] <= shark:
                    visited[nr][nc] = visited[r][c] + 1
                    queue.append((nr, nc))
    
    # result: 먹을 수 있는 먹이까지의 거리와 좌표 (종료 조건을 위해 큰 값을 초기화)
    result = [(inf, inf, inf)]

    # 먹을 수 있는 먹이 (상어 현재 크기보다 작음) 거리와 좌표 담아주기
    for r in range(N):
        for c in range(N):
            if 0 < arr[r][c] < shark:
                result.append((visited[r][c], r, c))

    # 거리가 최소이면서 좌표가 왼쪽 위인 먹이의 정보 반환
    return min(result)


# N: 배열 크기, arr: 배열 정보
N = int(input())
arr = [list(map(int, input().split())) for _ in range(N)]

# 아기 상어 위치 찾아주기 및 탐색을 위한 좌표 초기화
for r in range(N):
    for c in range(N):
        if arr[r][c] == 9:
            shark_r, shark_c = r, c
            arr[r][c] = 0

inf = float('INF')
# shark: 상어 현재 크기, size_cnt: 먹은 먹이 개수
shark = 2
size_cnt = 0

# ans: 상어 이동 시간
ans = 0

while True:
    # 다음 먹이 정보 찾아주기
    time, next_r, next_c = bfs(shark_r, shark_c)

    # 먹이가 없으면 중단
    if time == inf:
        break
    
    # 먹이 먹기 및 해당 좌표로 상어 이동
    arr[next_r][next_c] = 0
    shark_r, shark_c = next_r, next_c

    # 정답 더해주기
    ans += time

    # 먹이 먹은 개수 세주기
    size_cnt += 1

    # 상어 크기 커지는 경우
    if size_cnt == shark:
        shark += 1
        size_cnt = 0

# 정답 출력
print(ans)

* 느린 풀이

from collections import deque
import sys, os, io, atexit

input = lambda: sys.stdin.readline().rstrip('\r\n')
stdout = io.BytesIO()
sys.stdout.write = lambda s: stdout.write(s.encode("ascii"))
atexit.register(lambda: os.write(1, stdout.getvalue()))

# 221012 16236 아기상어

# 정답코드


def bfs(r, c):
    visited = [[0 for _ in range(N)] for _ in range(N)]

    visited[r][c] = 1

    queue = deque()
    queue.append((r, c))

    while queue:
        r, c = queue.popleft()

        for i in range(4):
            nr = r + dr[i]
            nc = c + dc[i]
            if 0 <= nr < N and 0 <= nc < N:
                if visited[nr][nc] == 0:
                    if arr[nr][nc] == 9:
                        return visited[r][c]

                    if arr[nr][nc] <= shark:
                        visited[nr][nc] = visited[r][c] + 1
                        queue.append((nr, nc))

    return float('inf')



N = int(input())
arr = [list(map(int, input().split())) for _ in range(N)]

for r in range(N):
    for c in range(N):
        if arr[r][c] == 9:
            shark_r, shark_c = r, c

shark = 2
size_cnt = 0


dr = [-1, 1, 0, 0]
dc = [0, 0, -1, 1]

ans = 0

while True:
    distance = float('inf')

    for r in range(N):
        for c in range(N):
            if 0 < arr[r][c] < shark:
                if bfs(r, c) < distance:
                    next_r, next_c = r, c
                    distance = bfs(r, c)

    if distance == float('inf'):
        break
    
    arr[shark_r][shark_c] = 0
    arr[next_r][next_c] = 9

    shark_r, shark_c = next_r, next_c

    ans += distance

    size_cnt += 1
    if size_cnt == shark:
        shark += 1
        size_cnt = 0

print(ans)