백준 1300 K번째 수

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

NxN의 행렬 A가 있을 때 A의 각 원소는 ixj로 표현된다고 하자. 이 때 A의 원소들을 오름차순으로 B에 저장했을 때 B[k]의 값을 구하는 문제이다. 이 때 인덱스는 1부터 시작한다고 한다.

다음은 n=4 일때를 예시로 A를 그려보았다.

• A행렬 (n = 4)

| 1 | 2 | 3 | 4 | | :–: | :–: | :–: | :–: | | 2 | 4 | 6 | 8 | | 3 | 6 | 9 | 12 | | 4 | 8 | 12 | 16 |

var end = k+1
var start = 1

한가지 포인트는 우리가 k라는 index를 찾고자하면 B[k]의 저장된 값은 k보다 항상 작거나 같다는 것이다. 따라서 B[k]에 저장된 값을 우리는 이분탐색으로 찾을 것인데 초기 end값을 k+1로 주고 시작한다.

while start < end {
    let mid = (start+end)/2
    var count = 0
    for i in 1...n {
        let tmp = min(mid/i, n)
        if tmp == 0 {break} 
        //만약 0이 나왔다면 그 이후에 대해서도 0일것이므로 break
        count += tmp
    }
    if count < k {
        start = mid+1
    }
    else {
        end = mid
    }
}

위의 A행렬을 가로로 한줄씩 살펴보도록하자. 우리가 8이라는 수를 골랐다고하자. 그렇다면 2번째 줄의 8보다 작거나 같은 수의 갯수는 8/2 = 4 가 될 것이다. 그렇다면 3번째줄의 경우는? 8/3 = 2 로 총 2개가 나타나게 된다. 즉 임의의 i번째 row의 k보다 작거나 같은 숫자의 갯수는 min(k/i, n)이 된다.

따라서 B[k]에 저장된 값을 이분탐색으로 찾을 때 mid보다 작거나 같은 원소의 갯수를 count라는 변수에 저장해준다. 해당 갯수가 k보다 작다면 B[k]의 값을 증가시켜야 하므로 우측 영역을 탐색하도록 하고 k보다 크거나 같다면 좌측 영역을 탐색하도록 하였다.

• 전체 코드

let n = Int(readLine()!)!
let k = Int(readLine()!)!
var end = k+1
var start = 1
while start < end {
    let mid = (start+end)/2
    var count = 0
    for i in 1...n {
        let tmp = min(mid/i, n)
        if tmp == 0 {break}
        count += tmp
    }
    if count < k {
        start = mid+1
    }
    else {
        end = mid
    }
}
print(end)