이분탐색을 이용해서 MN개의 조합에서 k번째 수를 찾는 놀라운 방법

1. 문제

 

1300번: K번째 수 (acmicpc.net)

1300번: K번째 수

 

2. 풀이

 

k 제한이 커서 그런지 이전에 배운 우선순위 큐 방법으로 풀면 시간초과로 통과를 못한다

 

import heapq
from sys import stdin

n = int(stdin.readline())

k = int(stdin.readline())

q = []

for i in range(1,n+1):
    
    heapq.heappush(q,(i*1,i,1))

for _ in range(k-1):
    
    _,ind1,ind2 = heapq.heappop(q)

    if ind2+1 < n+1:
        
        heapq.heappush(q,(ind1*(ind2+1),ind1,ind2+1))

value,_,_ = heapq.heappop(q)

print(value)

 

 

풀이를 잘 설명해준 글이 있어서 하나하나 따라가보자

 

 

 

[백준] 1300번 : K번째 수 - JAVA [자바] (tistory.com)

[백준] 1300번 : K번째 수 - JAVA [자바]

 

 

 

이 문제는 N*N개의 수를 1차원으로 오름차순 나열했을때, k번째 수가 얼마인지를 구하는 문제이다.

 

 

인덱스가 1부터 시작하니까 B[k]가 무엇인지를 찾는 문제와 같다.

 

위 그림에서 11번째 수를 찾으라 한다면 정답은 B[11] = 8이다.

 

가장 쉬운 방법은 N*N개의 수를 모두 구해 오름차순 정렬하고 k번째 수를 찾는건데,

 

N의 제한이 10만이니 최대 100억개의 수가 나오므로 완전탐색으로는 시간안에 찾을 수가 없다

 

그러면 모든 수를 구하지 않고서도 k번째 수를 찾아야 한다는 소리인데 어떻게 가능할까?

 

B[11] = 8이 무엇을 의미하는지 생각해본다면..

 

1) N*N개의 수를 오름차순 나열했을때, 11번째 수가 8이다.

 

2) 바꿔말하면, 8보다 작거나 같은 수는 B배열에서 최소 11개가 있다.

 

일반적으로 오름차순으로 나열된 B배열에서 B[k] = x라고 한다면, "x보다 작거나 같은 수가 최소 k개는 있다"라는 말과 같다.

 

오름차순으로 배열되었기 때문에 당연한 소리지만 막상 생각하기 쉽지 않다

 

그러면 이제 B[k] = x를 만족하는 어떤 x를 찾으면 된다는 소리

 

x를 바꿔나가면서, x보다 작거나 같은 수가 k개가 되는 그러한 x를 찾으면 된다는 말이다.

 

----------------------------------------------------------------------------------------------------------------------------------------------------------------

 

그러면 문제는 x보다 작거나 같은 수의 개수를 도대체 어떻게 세야하는가? 

 

아무리 생각해봐도 N*N개의 수를 구해야하는거 아닌가?

 

의외로 간단하면서도 놀라운 방법이 있다

 

배열 [1,2,3,4,5,6,7,8,9,10]에서 7보다 작거나 같은 수의 개수는 어떻게 셀 수 있을까? 1,2,3,4,5,6,7로 7개다.

 

그러면 이 배열의 원소에 2를 곱해서 [2,4,6,8,10,12,14,16,18,20]에서 7보다 작거나 같은 수의 개수는 어떻게 셀 수 있을까?

 

2,4,6 3개이다. 이거는 7을 2로 나눈 몫이다.

 

그러면 3을 곱해서 만든 [3,6,9,12,15,18,21,24,27,30]에서 7보다 작거나 같은 수의 개수는? 3,6 2개다. 이거는 7을 3으로 나눈 몫이다.

 

4를 곱해서 만든 [4,8,12,16,20,24,28,32,36,40]에서 7보다 작거나 같은 수의 개수는? 4로 1개이다. 이거는 7을 4로 나눈 몫이다.

 

...

 

정리하자면 다음과 같다.

 

1부터 N까지 있는 배열에서 k보다 작거나 같은 수의 개수는 k//1개

 

원소에 2를 곱한 2,4,6,8,10,...,2N까지 있는 배열에서 k보다 작거나 같은 수의 개수는 k//2개

 

..

 

원소에 i를 곱한 i,2i,3i,...,iN에서 k보다 작거나 같은 수의 개수는 k//i개

 

따라서 N*N 행렬 내에서 각 행 1행,2행,...,N행 각각에서 어떤 수 k보다 작거나 같은 수의 개수는 k//1,k//2, ..., k//N개이고

 

전부 더하면 N*N개의 수에서 k보다 작거나 같은 수의 개수를 구하지 않고도 셀수 있게 된다.

 

-------------------------------------------------------------------------------------------------------------------------------------------

 

따라서 가능한 x의 범위는 1부터 $n^{2}$까지 일거고, 이분탐색으로 x보다 작거나 같은 수의 개수가 최소 k가 되는 x를 찾으면 된다.

 

작거나 같은 수의 개수를 세는 방법은 1행,2행,...,n행에 대해 mid//i를 누적합 해주면 될거고

 

이분탐색 방법으로는 lower bound를 적용하면 될 것이다.

 

mid지점보다 작거나 같은 수의 개수가 k개인 경우가 여러가지 일 수 있다.

 

어떻게 여러가지일 수 있냐고? mid값이 B배열에 존재하는 수인 경우일 수도 있고 B배열에 존재하지 않는 수인 경우일 수도 있어서

 

 

위 경우 x = 5인 경우에 k = 8이고, x = 4인 경우에도  k = 8이다. 그렇지만 정답은 k = 8이 최초로 나오는 x = 4이고, 이것이 lower bound의 정의?였지

 

이분탐색 하도 안쓰다보니 구현 방법이 기억이 안나더라

 

이분탐색 연습좀 많이 해야겠어

 

 

이진 탐색 정복기3 -심화 응용 lower bound, upper bound- (tistory.com)

이진 탐색 정복기3 -심화 응용 lower bound, upper bound-

 

위에서 배운 알고리즘대로 lower bound를 구현해보면

 

def lower_bound(target,start,end):
    
    while start < end:
        
        mid = (start+end)//2

        sum = 0
        
        #A의 각 행에서 mid보다 작거나 같은 수의 개수를 세준다.
        for i in range(1,n+1):
            
            sum += mid//i
        
        #크면 클수록, 해당 수보다 작거나 같은 수의 개수는 증가하니까...
        
        #mid보다 작거나 같은 수의 개수가 target값보다 크거나 같다면, 
        #찾고자 하는 값은 mid보다 작거나 같다.
        if sum >= target:
            
            end = mid
        
        #mid보다 작거나 같은 수의 개수가 target보다 작다면..
        #찾고자 하는 값은 mid보다 클것이다.
        else:
            
            start = mid+1
        
    return end

 

 

이렇게 작성하면 마지막으로 문제점이 있는데...

 

mid보다 작거나 같은 수의 개수를 셀때, 우리는 A행렬에서 1행,2행,3행,...,N행 각각에서 작거나 같은 수의 개수를 세고 있는데

 

각 행에서 특정 수 k보다 작거나 같은 수의 개수는 최대 몇개가 있을까?

 

최대 N개다. 각 행의 길이가 N개니까

 

하지만 x의 범위가 1부터 N*N까지 가능한 상황에서 mid값이 N보다 클수도 있다.

 

그래서 단순히 sum += mid//i를 해버리면 N보다 큰 값이 더해질 수도 있다

 

def lower_bound(target,start,end):
    
    while start < end:
        
        mid = (start+end)//2

        sum = 0
        
        #A의 각 행에서 mid보다 작거나 같은 수의 개수를 세준다.
        for i in range(1,n+1):
            
            #길이가 n인 각 행에서 mid보다 작거나 같은 수의 개수는 최대 n개이다.
            sum += min(mid//i, n)
        
        #크면 클수록, 해당 수보다 작거나 같은 수의 개수는 증가하니까...
        
        #mid보다 작거나 같은 수의 개수가 target값보다 크거나 같다면, 
        #찾고자 하는 값은 mid보다 작거나 같다.
        if sum >= target:
            
            end = mid
        
        #mid보다 작거나 같은 수의 개수가 target보다 작다면..
        #찾고자 하는 값은 mid보다 클것이다.
        else:
            
            start = mid+1
        
    return end

 

 

따라서 n,k를 입력받고 lower_bound(k,1,n*n)을 구하면 정답이 나올건데

 

x의 범위를 1부터 n*n까지라고 했지만, 사실 더 줄일수 있다고 한다.

 

 

x가 k보다 클수가 없는게, 1부터 n까지 배열과 1부터 n까지 배열의 조합에서 항상 1씩 증가하는게 아니고

 

0이 증가하는, 중복되는 수가 반드시 존재하지만(B배열의 하얀색부분)

 

B배열의 인덱스는 중복되는 수 없이 1부터 증가하잖아(검정색부분)

 

하얀색 부분이 x이고 검정색 부분이 k인데, 그러므로 x <= k

 

억지같은데??;; 그래도 합리적이긴해

 

아무튼 1부터 n*n까지 탐색할 필요 없고 1부터 k까지 lower bound로 탐색하면 되겠다

 

from sys import stdin

def lower_bound(target,start,end):
    
    while start < end:
        
        mid = (start+end)//2

        sum = 0

        for i in range(1,n+1):
            
            sum += min(mid//i, n)

        if sum >= target:
            
            end = mid
        
        else:
            
            start = mid+1
        
   
    return end
            
n = int(stdin.readline())
k = int(stdin.readline())

print(lower_bound(k,1,k))