[백준/16236] 아기상어

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

16236번: 아기 상어

 

 삼성전자 공채 코딩테스트에 나왔던 기출문제입니다. 사실 이 문제를 풀었다는 것은 저에게 나름 뜻깊은 일입니다. 상반기 코딩테스트에서 좌절한 뒤, 처음 코딩테스트 알고리즘 공부를 시작했을 때 접했던 것이 이 문제입니다. 당시에는 정말 아무것도 몰랐기 때문에 '내가 정말 코딩테스트에 합격할 수 있을까?'라는 의문이 들 정도로 어렵게만 느껴졌습니다. 그러나 매일의 힘을 믿고 하루하루 정진한 결과 어느날 이 문제를 풀고야 말았네요. 자부심이 느껴지는 일입니다:)

 

---

풀이

기본적으로 DFS를 이용하여 접근하되, 다익스트라 알고리즘처럼 우선순위 큐를 이용해야 합니다. 소스코드는 다음과 같습니다.

import heapq


n = int(input())
graph = []
for _ in range(n):
    graph.append(list(map(int, input().split())))

level = 2
cnt = 0
time = 0

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

for i in range(n):
    for j in range(n):
        if graph[i][j] == 9:
            start_x, start_y = i, j
            graph[i][j] = 0
            break

while True:      
    visited = [[-1] * n for _ in range(n)]

    visited[start_x][start_y] = 0
    q = []
    heapq.heappush(q, (0, start_x, start_y))
    
    while q:
        d, x, y = heapq.heappop(q)
        if 0 < graph[x][y] < level:
            time += visited[x][y]
            start_x, start_y = x, y
            graph[x][y] = 0
            cnt += 1
            if cnt == level:
                level += 1
                cnt = 0
            break
        for i in range(4):
            nx, ny = x + dx[i], y + dy[i]
            if not (0 <= nx < n and 0 <= ny < n):
                continue
            if graph[nx][ny] > level:
                continue
            if visited[nx][ny] != -1:
                continue
            heapq.heappush(q, (visited[x][y] + 1, nx, ny))
            visited[nx][ny] = visited[x][y] + 1
    else:
        print(time)
        break

---

 입력을 받은 뒤 우선 아기상어의 위치를 찾아야 하기 때문에 그래프를 순회하며 9를 찾았습니다. 해당 위치는 start_x와 start_y로 따로 저장하고, 실제 graph상에서는 0으로 지워버립니다. 이는 나중에 graph에서 물고기 정보와 헷갈릴 수 있기 때문입니다. level이 만약 9를 넘어갈 경우 자신의 위치를 먹어버리는 일이 발생할 수 있습니다.

for i in range(n):
    for j in range(n):
        if graph[i][j] == 9:
            start_x, start_y = i, j
            graph[i][j] = 0
            break

   

 그 후 본격적으로 정보를 초기화시킵니다. visited 변수는 전체를 -1로 초기화한 뒤 시작점에서부터의 거리 정보를 담을 것입니다. 따라서 visited변수가 -1인지 아닌지를 기준으로 방문 여부를 판별할 수 있으며, 또한 해당 노드까지의 거리도 구할 수 있습니다.

 이 때 중요한 것은 우선순위 큐를 사용한다는 것입니다. 문제의 조건에서 '거리가 가까운 물고기부터 먹고,거리가 가까운 물고기가 많다면 가장 위의 물고기를 먹고, 가장 위의 물고기가 많다면 가장 왼쪽의 물고기부터 먹는다'고 제시하였습니다. 따라서 우선순위 큐의 최소큐 개념을 이용하여 (거리, x좌표, y좌표) 순으로 원소를 삽입하면 조건에 맞는 노드부터 뽑을 수 있습니다.

visited = [[-1] * n for _ in range(n)]

visited[start_x][start_y] = 0
q = []
heapq.heappush(q, (0, start_x, start_y))

 

 다음은 본격적으로 DFS를 수행하는 구간입니다. 우선순위 큐에서 원소를 뽑아 만약 해당 노드의 값이 0보다 크고(물고기가 있고) level보다 작다면(먹을 수 있다면) 총 흐른 시간을 저장할 time변수에 해당 노드까지의 거리를 더해줍니다. 또한 새로운 순회를 위해서 start_x와 start_y를 도착지점으로 초기화해주고, 앞선 경우와 마찬가지로 그래프 내에서 0으로 지워줍니다. 레벨만큼 물고기를 먹으면 레벨업하는 형식이기 때문에 따로 cnt변수를 이용하여 물고기를 먹을 때마다 1씩 더해줍니다. 만약 cnt변수가 level과 같아지면 레벨업시켜주고, cnt는 다시 0으로 초기화합니다.

 이 조건문에 걸리지 않는다면 상하좌우를 확인합니다. 단, 1) 좌표를 벗어나는 경우 2) 크기가 더 큰 물고기가 있어서 지나가지 못하는 경우 3) 이미 방문한 경우에 대해서는 고려하지 않도록 해줍니다. 예외사항에 걸리지 않는다면 (새로운 거리, 새로운 x좌표, 새로운 y좌표) 형태로 우선순위 큐에 삽입하고 마찬가지로 visited 변수에 거리도 업데이트해줍니다.

 만약 q를 다 돌았는데 break가 걸리지 않는다면 더 이상 먹을 수 있는 물고기가 없다는 의미이므로,  이때는 전체 흐른 시간인 time을 출력해줍니다. 브레이크 체커인 else가 곳곳에서 정말 요긴하게 쓰이는 듯 합니다 :)

while q:
        d, x, y = heapq.heappop(q)
        if 0 < graph[x][y] < level:
            time += visited[x][y]
            start_x, start_y = x, y
            graph[x][y] = 0
            cnt += 1
            if cnt == level:
                level += 1
                cnt = 0
            break
        for i in range(4):
            nx, ny = x + dx[i], y + dy[i]
            if not (0 <= nx < n and 0 <= ny < n):
                continue
            if graph[nx][ny] > level:
                continue
            if visited[nx][ny] != -1:
                continue
            heapq.heappush(q, (visited[x][y] + 1, nx, ny))
            visited[nx][ny] = visited[x][y] + 1
    else:
        print(time)
        break

 위와 같은 과정을 브레이크 체커에 의해서 순회가 종료될 때까지 반복해주면 정답을 구할 수 있습니다.