[이진탐색/백준] 1300번 : K번째 수

- 문제 : https://www.acmicpc.net/problem/1300

 

본 풀이는 인프런의 '떠먹는 알고리즘 코딩테스트 with 파이썬'을 참고해 작성했습니다. 

 

문제 이해 

n x n 행렬에서 각 원소는 i x j 값으로 이루어져있고, 전체 원소들을 정렬 시킨 뒤 k번째 수를 찾는 문제였다. 

직관적으로는 n x n 행렬을 만들어서 원소를 추가한 뒤 모든 수를 정렬해 k번째 원소를 찾으면되는 쉬운 문제라고 생각했다. 

하지만 위처럼 진행하면 공간복잡도 O(N^2), 시간 복잡도(N^2logN)으로 메모리/시간 초과가 분명 발생하고, n의 범위 또한 최대 10만이라 성공하기 어려운 로직이었다. 

 

그래서 이 문제를 공간복잡도, 시간복잡도가 효율적이게 풀려면 다음과 같은 아이디어가 필요하다.

N개의 원소 중에서 k번째 수는 다음과 같은 조건을 만족한다. 

- k번째 수보다 작은 수는 총 (k - 1)개 이하다

- k번째 수보다 큰 수는 총 (n - k)개 이하다. (n은 전체 배열의 크기)

 

풀이 로직

위 아이디어를 바탕으로 풀이 로직을 정리하자면 다음과 같다. 

1. k번째 숫자가 될 후보 숫자들을 이진탐색으로 구한다. 

2. 가장 왼쪽 값(low)은 1, 가장 오른쪽 값(high)은 n * n으로 시작한다.

3. 각 mid 값을 기준으로 1행부터 n행까지 반복문을 돌면서 mid 보다 작은 수의 개수와 큰 개수를 합하고, 위의 아이디어를 만족하는지 확인하면 된다. 

3-1. 만약 작은 수가 많다면 더 작은 숫자들이 배열을 구성하게 만들고 (high = mid -1)

3-2. 만약 큰 수가 많다면 더 큰 숫자들이 배열을 구성하게 만들면 된다. (low = mid + 1)

 

여기에서 3번의 mid보다 작은 수/큰 수 개수를 어떻게 구해야할지 잘 몰랐다. 

mid보다 작은 수의 개수를 세는 것을 기준으로 이해하면 다음과 같다. 

위의 로직으로 i행에서 x값보다 작은 값을 가지는 값은 1열부터 min(n, (x - 1) // i)라는 것을 알 수 있다. 

같은 i행에서 큰 수는 i * j 값을 포함한 x보다 작은 값, min(n, x  // i)를 n에서 빼주면 된다. 

 

 

 

- 정답 풀이 

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

def get_num_smaller(x):
    num_smaller = 0
    for i in range(1, n + 1):
        num_smaller += min(n, (x - 1) // i)
        
    return num_smaller

def get_num_bigger(x):
    num_bigger = 0
    for i in range(1, n + 1):
        num_bigger += n - min(n, x // i)
        
    return num_bigger

low, high = 1, n * n
answer = -1 
while low <= high:
    mid = (low + high) // 2
    
    num_smaller = get_num_smaller(mid)
    num_bigger = get_num_bigger(mid)
    
    if num_smaller > k - 1:
        high = mid - 1
    elif num_bigger > n * n - k:
        low = mid + 1
    else:
        answer = mid 
        break
        
print(answer)