[알고리즘] BOJ 1300 K번째 수 #Python

[BOJ] 1300 K번째 수

📍 문제

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

배열 A와 B의 인덱스는 1부터 시작한다.

---

📍 입력

첫째 줄에 배열의 크기 N이 주어진다. N은 105보다 작거나 같은 자연수이다. 둘째 줄에 k가 주어진다. k는 min(109, N2)보다 작거나 같은 자연수이다.

---

📍 출력

B[k]를 출력한다.

---

📍 풀이

• 이분 탐색(binary search) 알고리즘을 이용하였다.

• A[i][j] = i×j 를 원소로 갖는 N×N인 배열 A 에서 어떤 수(mid) 보다 작은 값의 개수(cnt)를 구하는 방법은 1 부터 N 까지의 수를 mid 의 값으로 나눈 몫과 N 의 값 중 더 작은 값을 모두 더한 값과 같다.

  N = 4, mid = 5 일 경우, 배열 A 에서 4보다 작은 값의 개수는
  
  [1, 2, 3, 4] > min(4, 5//1) = 4
  [2, 4, 6, 8] > min(4, 5//2) = 2
  [3, 6, 9, 12] > min(4, 5//3) = 1
  [4, 8, 12, 16] > min(4, 5//4) = 1
  
  로 총 8이다.

  • 위의 방법을 이용하여 mid 값보다 작은 값의 개수(cnt)가 K 보다 작다면 start 의 값에 mid + 1 값을 할당하고,
  • 반대로 mid 값보다 작은 값의 개수(cnt)가 K 보다 크다면 end 의 값에 mid - 1 의 값을 할당한다.

• 이분 탐색을 통해 K 번째 인덱스에 존재할 수 있는 값을 반환한다.

✍ 코드

from sys import stdin

def func(N,K):
    answer = 0
    start, end = 1, N**2
    while start <= end:
        mid = (start + end) // 2
        cnt = 0
        for i in range(1,N+1):
            cnt += min(N,mid//i)
        if cnt < K:
        	start = mid + 1
        else:
            answer = mid
            end = mid - 1
    return answer

N = int(stdin.readline())
K = int(stdin.readline())
print(func(N,K))