[고득점 Kit (이분탐색)] 징검다리 (Swift, Python)

코딩테스트 연습 - 징검다리

 

풀이

 바위가 최대 50,000개까지 있을 수 있고 이 중에서 n개를 제거해야 하는데 n개를 제거하는 경우의 수를 따지면 시간초과가 날 것이다. 이분탐색을 사용하는데 어떤 값을 기준으로 이분탐색을 해야할지 정해야한다. 문제에서 바위를 n개 제거한 뒤 각 지점 사이의 거리의 최솟값 중에 가장 큰 값을 구하라고 하였다. 그렇기 때문에 각 지점 사이의 거리를 기준으로 하여 이분탐색을 진행했다.

 

 시작점은 0으로 두고 끝점은 입력 값으로 주어진 distance로 두었다. 바위를 모두 제거할 수 있을 때 25가 거리의 최솟값이기 때문이다. 그 후 mid를 구하여 출발지점부터 mid만큼 거리를 이동했을 때 바위의 수를 카운트한다. 문제의 예제로 알아보자. 바위가 예제처럼 2, 11, 14, 17, 21이 있다고 하자. 현재 mid가 12라고 하면 시작지점인 0에서 12만큼 이동한다. 그럼 다음으로 존재할 수 있는 바위는 14가 된다. 즉, 2와 11에 있는 바위가 제거된 상태이다. 다음에는 14에서 mid인 12를 더한다. 그럼 26이 된다. 해당 위치 이후에는 바위가 존재할 수 없다. 즉, 17과 21에 있는 바위가 제거된 상태이다.

 이렇게 될 경우 12는 답이 될 수 없다. 2개의 바위만 제거해야 하는데 4개의 바위가 제거되었기 때문이다. 그러면 끝점을 11로 설정하고 이분탐색을 진행한다.

 

 위와 같은 방법으로 이분탐색을 진행해 나가면 문제에서 요구하는 값을 구할 수 있다.

 

코드

[Swift]

import Foundation

func solution(_ distance:Int, _ rocks:[Int], _ n:Int) -> Int {
    let sortedRocks: [Int] = rocks.sorted() // 먼저 입력으로 주어진 바위의 위치 정렬
    var answer: Int = 0
    
    var left: Int = 0
    var right: Int = distance
    
    while left <= right {
        let mid: Int = (left + right) / 2 // mid가 각 지점 사이의 거리의 최솟값
        var cnt: Int = 0 // mid를 기준으로 했을 때 남아있는 바위의 개수
        var now: Int = mid // now는 현재 위치. 0부터 시작해서 mid만큼 거리가 떨어져 있으므로 처음에는 mid로 설정해준다.
        
        // 이제 바위의 위치를 for문으로 돌면서 남아 있는 바위의 개수를 카운트한다.
        // 현재 바위의 위치가 now 이상이면 해당 바위는 제거되지 않은 바위가 된다.
        // 바위의 개수를 1 증가시키고 현재 위치 now를 남은 바위의 위치(rock)에서 mid만큼 이동시켜준다.
        for rock in sortedRocks {
            if rock >= now {
                cnt += 1
                now = rock + mid
            }
        }
        
        // 남아있는 바위의 개수가 n개를 제거한 이후의 바위의 개수 이상일 경우가 각 지점 사이의 거리의 최솟값 후보가 된다.
        if cnt >= rocks.count - n {
            left = mid + 1
            answer = mid
        } else {
            right = mid - 1
        }
    }
    
    return answer
}

[Python]

def solution(distance, rocks, n):
    rocks.sort() # 먼저 입력으로 주어진 바위의 위치 정렬
    answer = 0 

    left = 0
    right = distance

    while left <= right:
        mid = (left + right) // 2  # mid가 각 지점 사이의 거리의 최솟값
        cnt = 0  # mid를 기준으로 했을 때 남아있는 바위의 개수
        now = mid # now는 현재 위치. 0부터 시작해서 mid만큼 거리가 떨어져 있으므로 처음에는 mid로 설정해준다.

        # 이제 바위의 위치를 for문으로 돌면서 남아 있는 바위의 개수를 카운트한다.
        # 현재 바위의 위치가 now 이상이면 해당 바위는 제거되지 않은 바위가 된다.
        # 바위의 개수를 1 증가시키고 현재 위치 now를 남은 바위의 위치(rock)에서 mid만큼 이동시켜준다.
        for rock in rocks:
            if rock >= now:
                cnt += 1
                now = rock + mid
        
        # 남아있는 바위의 개수가 n개를 제거한 이후의 바위의 개수 이상일 경우가 각 지점 사이의 거리의 최솟값 후보가 된다.
        if cnt >= len(rocks) - n:
            left = mid + 1
            answer = mid
        else:
            right = mid - 1

    return answer