[Python] K번째 수 - 백준 (이진탐색)

해당 글은 백준 1300번 'K번째 수' 문제를 다룬다.

문제

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

1300번: K번째 수

 

설명

해당 문제는 2중 for문으로 행렬을 만들고, 인덱스로 접근했을땐 메모리 초과가 떠서 해당 방법으론 풀 수 없었다.

이진탐색으로 nxn 행렬에서 어떤 수보다 작거나 같은 값이 몇 개인지 찾아야 몇 번째 숫자인지 알 수 있다.

 

예를 들어, 3x3 행렬의 값들을 리스트로 정렬해보면 [1, 2, 2, 3, 3, 4, 6, 6, 9]와 같다.

이때 5보다 작거나 같은 값을 찾아보면, 아래의 수가 존재한다. 

1행 : 1x1 ~ 1x3

2행 : 2x1 ~ 2x2

3행 : 3x1

----------- 이는 5를 행으로 나눈 몫과 같다. 이때 나눈 몫의 값이 열의 수를 초과할 수 없으므로 min(값//행, n)이 된다.

1행 : 5//1 -> min(5, n) -> min(5, 3) -> 3개

2행 : 5//2 -> 2개

3행 : 5//3 -> 1개

즉, 5보다 작거나 같은 값은 6개가 된다.

이런식으로 1~nxn의 수에서 이진탐색으로 mid 값보다 작거나 같은 수를 찾아가면 mid가 몇 번째 숫자인지 알 수 있다.

 

코드

n = int(input())
k = int(input())

start, end = 1, n*n
result = 0
while start <= end:
    mid = (start+end)//2
    count = 0
    for i in range(1, n+1):
        count += min(mid//i, n)
    if count >= k:
        result = mid
        end = mid - 1
    else:
        start = mid + 1
print(result)