[BOJ / Swift & Python] 16236 - 아기 상어

BOJ 16236번 아기 상어

16236번: 아기 상어

* 한 게시글에 코드 블록 언어를 하나로 통일해야 해서 Python 코드는 하이라이트가 정상적이지 않습니다ㅠ

Intro

처음에는 물고기들의 위치를 저장한 뒤, 각 물고기의 위치에서 현재 상어 위치까지의 거리를 구해 우선순위에서 일 순위인 물고기를 사용하는 방식으로 구현하였습니다. 그랬더니 시간 초과가 계속해서 발생했습니다.

위 풀이 방식은 가장 가까우면서 우선순위에서 일순위인 물고기를 찾기 위해, BFS를 상어가 먹을 수 있는 물고기 개수만큼 돌려야 합니다.

하지만 가장 가까우면서 우선순위에서 일순위인 물고기를 찾는 것은 한 번의 BFS면 충분했습니다. 그 풀이 방법을 공유해보도록 하겠습니다.

Solution

1. 입력을 받으면서 숫자 9가 입력되면 상어의 위치를 저장하고, 9 숫자를 0으로 변경합니다.
2. 상어의 현재 위치에서 가장 가까운 거리의 물고기를 BFS로 찾아 물고기의 위치와 상어와 물고기와의 거리를 반환합니다.
   • BFS
     가장 가까운 거리의 물고기들을 우선순위에 따라 오름차순으로 정렬한 뒤, 가장 첫 번째 값을 반환합니다.
     먹을 수 있는 물고기가 없을 경우 불가능한 값(위치를 (-1,-1)로 하는 등)을 반환합니다.
     • 우선순위
       1. 상어와의 거리
       2. 위에서부터 떨어진 거리
       3. 왼쪽에서부터 떨어진 거리
3. 입력값에서 반환된 물고기 위치의 값을 0으로 바꾸고,
   거리만큼 time을 증가시키고, 먹은 물고기 개수를 1만큼 증가시킵니다.
4. 먹은 물고기 개수가 상어의 크기와 동일하면, 상어의 크기를 1만큼 증가시키고 먹은 물고기 개수를 0으로 초기화합니다.
5. 더 이상 먹을 수 있는 물고기가 없을 때(2번의 반환 값이 불가능한 값일 때)까지 2-4번을 반복합니다.

Code

Swift

import Foundation

let n = Int(readLine()!)!
var space = Array<[Int]>()
var sharkPos = (x: -1, y: -1)
var sharkSize = 2
for x in 0..<n {
    space.append(readLine()!.split(separator: " ").map(){Int(String($0))!})
    for (y, num) in space.last!.enumerated() where num == 9 {
        space[x][y] = 0
        sharkPos = (x: x, y: y)
    }
}

func bfs(x: Int, y: Int, n: Int, sharkSize: Int, space: [[Int]]) -> (pos: (x: Int, y: Int), time: Int) {
    var queue = [[x, y, 0]]
    var visited = Array<[Bool]>(repeating: Array(repeating: true, count: n), count: n)
    var fish = [[-1, -1, n*n+1]]
    while !queue.isEmpty {
        let a = queue.removeFirst()
        if fish.count > 1 && fish.last![2] < a[2] { break } 
        else if 1..<sharkSize ~= space[a[0]][a[1]] { fish.append(a); continue }
        else if !visited[a[0]][a[1]] || space[a[0]][a[1]] > sharkSize { continue }
        else { visited[a[0]][a[1]] = false }
        if a[0] != 0 { queue.append([a[0]-1, a[1], a[2]+1]) }
        if a[1] != 0 { queue.append([a[0], a[1]-1, a[2]+1]) }
        if a[1] != n-1 { queue.append([a[0], a[1]+1, a[2]+1]) }
        if a[0] != n-1 { queue.append([a[0]+1, a[1], a[2]+1]) }
    }
    fish.sort(){ (l, r) -> Bool in
        if l[2] < r[2] { return true }
        else if l[2] == r[2] {
            if l[0] < r[0] { return true }
            else if l[0] == r[0] {
                return l[1] < r[1]
            }
            else { return false }
        }
        else { return false }
    }
    return ((fish[0][0], fish[0][1]), fish[0][2])
}

var time = 0
var fishCount = 0
while true {
    let result = bfs(x: sharkPos.x, y: sharkPos.y, n: n, sharkSize: sharkSize, space: space)
    if result.pos.x < 0 { break }
    time += result.time
    sharkPos = result.pos
    fishCount += 1
    if fishCount == sharkSize {
        sharkSize += 1
        fishCount = 0
    }
    space[sharkPos.x][sharkPos.y] = 0
}
print(time)

Python

import sys
input = sys.stdin.readline

babyshark = 2
n = int(input().strip())
space = []
babySharkPos = (0, 0)
for x in range(n) :
    space.append(list(map(int, input().strip().split())))
    for y, f in enumerate(space[-1]) : 
        if f == 9 : 
            babySharkPos = (x, y)
            space[-1][y] = 0

def distance(start):
    q = [(start, 0)]
    visited = [[True]*n for _ in range(n)]
    fish = [[(-1, -1), n*n+1]]
    while len(q) :
        a = q.pop(0)
        if len(fish) > 1 and fish[-1][1] < a[1] : break 
        if 0 < space[a[0][0]][a[0][1]] < babyshark : fish.append(a)
        if not visited[a[0][0]][a[0][1]] or space[a[0][0]][a[0][1]] > babyshark : continue
        else : visited[a[0][0]][a[0][1]] = False
        if a[0][0] != 0 : q.append([(a[0][0]-1, a[0][1]), a[1]+1])
        if a[0][1] != 0 : q.append([(a[0][0], a[0][1]-1), a[1]+1])
        if a[0][0] != n-1 : q.append([(a[0][0]+1, a[0][1]), a[1]+1])
        if a[0][1] != n-1 : q.append([(a[0][0], a[0][1]+1), a[1]+1])
    fish.sort(key=lambda x: (x[1], x[0][0], x[0][1]))
    return fish[0]
    
time = 0
fishCount = 0
while True :
    result = distance(babySharkPos)
    if result[0][0] < 0 : break
    time += result[1]
    babySharkPos = result[0]
    fishCount += 1
    if fishCount == babyshark :
        babyshark += 1
        fishCount = 0
    space[babySharkPos[0]][babySharkPos[1]] = 0
print(time)

* 개인 블로그에 작성한 내용을 복사해왔습니다.