BOJ #1300 K번째 수

읽기 전

• 불필요한 코드나 잘못 작성된 내용에 대한 지적은 언제나 환영합니다.
• 개인적으로 배운 점을 정리한 글입니다.

문제 링크

문제 1300. K번째 수

문제 풀이

이진 탐색을 적용해서 풀 수 있는 문제다. 이진 탐색으로 변형할 수 있는 문제들의 공통점을 유념할 필요가 있다.

• 탐색의 범위의 최솟값과 최댓값이 명확할 경우
• 전체 탐색 범위가 정렬된 형태이거나 $O(n)$ 이내에 문제에서 요구하는 수치를 계산할 수 있을 경우

결론적으로 문제 해결에 요구되는 함수의 형태가 단조함수인 경우를 의미한다.

명확하게 떨어지지 않겠지만 위 두 조건이 문제에 제시되었다면 브루트포스로 해결해야 하는 경우가 아니라면 문제의 해결에 요구되는 시간복잡도를 $O(nlog\ n)$으로 두고 탐색 횟수를 $O(log\ n)$, 정답 수치 계산을 $O(n)$으로 두고 해법을 찾아보자.

편하게 생각하면 탐색의 최솟값을 1, 최대값을 $N^2$으로 두고 탐색해도 되겠지만 최댓값을 K로 설정해도 무방하다. 그 이유는 아래와 같다.

탐색의 최댓값이 K여도 되는 이유

A를 정렬되지 않은 1차원 형태의 B' 배열로 만들었다고 치자. $B' = {1,\cdots, n, 2,\cdots, 2n, 3, \cdots, n^2 }$의 형태가 될 것이다. 그리고 K는 n보다 작거나 같거나 큰 3가지 경우의 수가 존재하는데 K가 n보다 작을 경우 K보다 작은 수는 K-1개가 된다. K가 n과 같을 땐 K보다 작은 수는 n-1이 되고 K가 N보다 큰 경우는 다시 상수가 곱해진 형태로 1~n을 반복하므로 중복된 숫자가 발생하여 K번 좌표의 수는 K보다 작을 수밖에 없다.

B[idx]보다 작은 A의 원소 개수

이진 탐색의 적용 방법은 최소 좌표를 1로 두고 최대 좌표를 K로 설정한 뒤 탐색 좌표에 대해 작거나 같은 A의 원소가 얼마나 있는지 계산해야 한다. 이 문제는 최소, 최대 좌표 설정은 어렵지 않으나 A의 원소를 카운트 하는 함수 작성이 난이도를 결정했다.

그렇다면 특정값 V에 대해 보다 작거나 같은 A의 원소의 개수를 계산하는 방법을 구상하자. 우선 idx에 대해 A의 각 행 1~n까지 나눈다고 했을 때 V // i가 V보다 작거나 같은 원소의 개수임은 알 수 있다. 그리고 V //i > N인 경우 행렬의 범위를 초과하기 때문에 min 함수로 보정해준다. 이런 식으로 n행에 대해 모두 수행하면 V보다 작거나 같은 A의 원소의 개수를 셀 수 있고 V에 대한 개수가 요구하는 좌표와 일치한다면 그 값이 B[idx]가 될 것이다.

python 코드

import sys
input = sys.stdin.readline

N = int(input())
k = int(input())

def search(N, k):
    def count_num(val):
        result = 0
        for i in range(1, N + 1):
            result += min(val // i, N)
        return result
    min_val,  max_val = 1, k
    sol = 0
    while min_val < max_val:
        mid_val = min_val + (max_val - min_val) // 2
        cnt = count_num(mid_val)
        if cnt >= k:  # cnt가 k보다 같거나 큰 수들 중 최소값
            max_val = mid_val - 1
            sol = mid_val
        elif cnt < k: # lower bound로 잡으면 k - 1 직후의 값을 잡는다.
            min_val = mid_val + 1
             # 그러면 배열에는 없으나 조건을 만족하는 수를 리턴해버림
    return sol

print(search(N, k))