1300. K번째 수 (Swift, Python)

1300번: K번째 수

 

풀이

 이분탐색을 활용해서 해결할 수 있다. B 배열에서 K번째에 있는 수의 범위는 1 이상 K 이하의 수이다. left를 1로 잡고 right를 K로 잡아서 이분 탐색을 진행하면 된다.

 

 mid를 구해서 해당 mid 이하인 원소 개수가 A 배열에서 몇 개인지를 구하면 된다. 예를 들어 현재 mid가 13이고 N이 4라고 가정해보자. A 배열에 13 이하인 원소의 개수는 1부터 N(4에 해당)까지 for문을 통해 구할 수 있다. 1행에서 4개, 2행에서 4개, 3행에서 4개, 4행에서 3개가 있다. 즉, A에는 13 이하인 원소의 개수가 총 15개 있다. 총 개수와 K를 비교했을 때 K가 더 크면 left를 늘려서 다시 탐색하고 그렇지 않으면 right를 줄여서 탐색한다. 

 

 

코드

[Swift]

var N: Int = 0
var K: Int = 0

if let input = readLine() {
    N = Int(input)!
}

if let input = readLine() {
    K = Int(input)!
}

var left: Int = 1
var right: Int = K

var res: Int = -1

while left <= right {
    let mid: Int = (left + right) / 2
    
    // A 배열에 mid 이하인 수가 몇 개 있는지 구한다.
    var cnt: Int = 0
    
    // 1행부터 N행까지 차례대로 구해준다.
    // mid와 i를 나눈 몫이 개수가 된다. 단, i행에는 N개의 원소만 있으므로 N 보다 몫이 더 크다면 N으로 설정한다.
    for i in 1...N {
        cnt += min(N, mid / i)
    }
    
    // mid 이하인 개수가 많으면 mid는 답이 될 수 있기 때문에 res에 mid를 설정하고
    // right를 줄여서 다시 이분탐색을 진행한다.
    // 그렇지 않을 경우 left를 늘려서 이분탐색을 진행한다.
    if cnt >= K {
        res = mid
        right = mid - 1
    } else {
        left = mid + 1
    }
}


print(res)

 

[Python]

if __name__ == '__main__':
    N = int(input())
    K = int(input())

    left = 1
    right = K

    res = -1

    while left <= right:
        mid = (left + right) // 2

        # A 배열에 mid 이하인 수가 몇 개 있는지 구한다.
        cnt = 0 

        # 1행부터 N행까지 차례대로 구해준다.
        # mid와 i를 나눈 몫이 개수가 된다. 단, i행에는 N개의 원소만 있으므로 N 보다 몫이 더 크다면 N으로 설정한다.
        for i in range(1, N+1):
            cnt += min(N, mid // i)

        # mid 이하인 개수가 많으면 mid는 답이 될 수 있기 때문에 res에 mid를 설정하고
        # right를 줄여서 다시 이분탐색을 진행한다.
        # 그렇지 않을 경우 left를 늘려서 이분탐색을 진행한다.
        if cnt >= K:
            res = mid
            right = mid - 1
        else:
            left = mid + 1

    print(res)