[백준] (1300) K번째 수 [Python]

참고 자료

• Do it! 알고리즘 코딩 테스트 - 파이썬 편

문제 링크

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

문제

세준이는 크기가 N×N인 배열 A를 만들었다. 배열에 들어있는 수 A[i][j] = i×j 이다. 이 수를 일차원 배열 B에 넣으면 B의 크기는 N×N이 된다. B를 오름차순 정렬했을 때, B[k]를 구해보자. 배열 A와 B의 인덱스는 1부터 시작한다.

• 입력
  첫째 줄에 배열의 크기 N이 주어진다. N은 105보다 작거나 같은 자연수이다. 둘째 줄에 k가 주어진다. k는 min(10⁹, N²)보다 작거나 같은 자연수이다.
• 출력
  B[k]를 출력한다.

예제 입력	예제 출력
3 7	6

---

문제 설명

k는 min(10⁹, N²)보다 작거나 같은 자연수이므로 시간 복잡도가 N²인 알고리즘은 사용할 수 없다. 대신 이진 탐색으로 절반씩 줄여가면서 B[k]를 구해야 한다. (사실 이 문제를 보고 이진 탐색을 떠올리기란 쉽지 않다.)

n = 3, k = 7일 때

최초의 mid 값은 4이다. 각 행에서 mid 보다 작거나 같은 수의 개수는 mid 값을 각 행의 첫번째 값으로 나눈 값이다. 그 결과 각 열에서 mid=4 이하인 수는 6개가 된다. mid 이하인 수의 개수는 min(mid // i, n)으로 계산한다. 그리고 이를 통해 mid=4는 6번째 수보다 큰 수가 될 수 없다는 것을 알 수 있으며, mid=4보다 큰 범위에 정답이 있다는 것을 유추할 수 있다.

 

mid보다 작은 수의 개수가 k보다 작으면 시작 인덱스를 mid+1로, k보다 크거나 같으면 종료 인덱스를 mid-1로 하면서 정답을 mid로 업데이트하며, 시작 인덱스가 종료 인덱스보다 커질 때까지 이진 탐색을 진행한다.

코드

import sys
input = sys.stdin.readline

n = int(input())
k = int(input())
s, e = 1, k
answer = 0

while s <= e:
    mid = (s+e)//2
    cnt = 0

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

    if cnt < k:
        s = mid + 1
    else:
        answer = mid
        e = mid - 1

print(answer)