백준 1300 K번째 수 Python 3

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

1300번: K번째 수

 

최종 코드

• GitHub -> https://github.com/hwayeon351/BEAKJOON-Algorithms/blob/main/BJ1300.py

GitHub - hwayeon351/BEAKJOON-Algorithms: 백준 알고리즘 소스 코드 모음

import sys

def find_kth_number(N, k):
    left = 1
    right = N*N
    while left <= right:
        mid = (left+right)//2
        cnt = 0
        for i in range(1, N+1):
            cnt += min(mid//i, N)
        if cnt >= k: right = mid - 1
        else: left = mid + 1
    return left

input = sys.stdin.readline
N = int(input())
k = int(input())
print(find_kth_number(N, k))

---

풀이 과정

NxN인 배열 A의 A[i][j] = ixj이다. 이 배열 안의 요소들을 오름차순으로 정렬했을 때, k번째 수가 몇인지를 구하는 문제이다.

먼저 단순하게 배열 A를 구하고 이를 오름차순 정렬을 한 다음 k번째 요소를 출력할 수 있지만, N은 1~10^5의 자연수이기 때문에 시간초과가 난다.

따라서 이분탐색을 적용해서 문제를 해결해야 한다.

그치만 어떻게 이분탐색을 적용하면 될지 감이 오지 않았다... 적당히 고민해보고 스스로 풀려고 하는 집념도 중요하지만 정말 모르겠다면 얼른 다른사람의 코드를 보고 내것으로 만드는 것도 중요하기 때문에 참고해서 문제를 풀었다.

 

이분탐색의 탐색 대상은 1~N^2 사이의 수이다. 배열 A에서 이 수보다 작거나 같은 수가 몇 개인지를 카운팅하면 해당 수가 몇 번째 수인지 알 수 있다.

탐색 대상을 mid로 하였을 때, mid 이하의 수가 배열 A에 몇 개가 있는지 알 수 있는 방법은 첫번째 row(또는 col)부터 마지막 row까지 mid//i를 한 값으로 알 수 있다.

이때 row(또는 col)은 1~N이기 때문에 min(mid//i, N)이 해당 row(또는 col)에 있는 수들 중 mid보다 작거나 같은 개수가 된다.

이렇게 각 row(또는 col)에 mid 이하의 숫자가 몇 개인지 구하고 모두 더하면 mid 이하의 수가 배열 A에 몇 개가 있는지 알 수 있다.

 

말로 정리한 내용은 이러한데 글만 봐서 이해가 되지 않기 때문에 하나의 예를 들어 구해보자.

배열 A가 4*4인 경우, 9번째 수는 6이다.

1	2	3	4
2	4	6	8
3	6	9	12
4	8	12	16

row는 1~4이기 때문에 1번 row에서 6 이하의 수는 1, 2, 3, 4로 min(6//1, 4) = 4개이다.

2번 row에서 6 이하의 수는 2, 4, 6으로 min(6//2, 4) = 3개이다.

3번 row에서 6 이하의 수는 3, 6으로 min(6//3, 4) = 2개이다.

4번 row에서 6 이하의 수는 4로 min(6//4, 4) = 1개이다.

따라서, 배열 A에서 6 이하의 수는 총 10개이다.

이때, 6은 배열 A에 2개 존재하므로 9번째 수가 될 수도 있지만 10번째 수가 될 수도 있다.

이는 앞이나 뒤의 수는 몇개 존재하는지를 통해서 6이 몇 번째 수인지 알 수 있다.(5 이하의 수는 8개, 7 이하의 수는 10개)

따라서, 이분 탐색을 적용해서 mid 숫자 이하의 수가 몇 개 있는지에 따라 left와 right 범위를 조정해서 K번째 수를 찾을 수 있다.

import sys

def find_kth_number(N, k):
	#가능한 숫자 범위 초기화
    left = 1
    right = N*N
    
    #이분 탐색
    while left <= right:
    	#mid를 탐색하고자 하는 숫자로 가정
        mid = (left+right)//2
        
        #mid 이하의 개수가 몇 개인지 카운팅
        cnt = 0
        for i in range(1, N+1):
            cnt += min(mid//i, N)
        
        #cnt가 k이상인 경우, right 범위를 줄인다
        if cnt >= k: right = mid - 1
        #cnt가 k미만인 경우, left 범위를 늘린다
        else: left = mid + 1
    return left

input = sys.stdin.readline
N = int(input())
k = int(input())
print(find_kth_number(N, k))

#시간복잡도 = O(log n), 공간복잡도 = O(1)