[Python] 백준 1300 - K번째 수

백준 1300 - K번 째 수 [GOLD 1]

두 개의 입력이 들어옵니다.

 

N과 K.

 

길이가 N인 정사각형 형태의 행렬이 arr이 있고, arr[i][j]에는 i*j가 존재합니다.

 

단, 인덱스는 1부터 시작합니다.

 

즉, N이 5일 때, 행렬은 다음과 같이 생성됩니다.

[[1, 2, 3, 4, 5],
 [2, 4, 6, 8, 10],
 [3, 6, 9, 12, 15],
 [4, 8, 12, 16, 20],
 [5, 10, 15, 20, 25]]

 

그리고 해당 2차원 행렬을 1차원 리스트로 풀어내고, 오름차순 정렬을 시킨 뒤, [K] 값을 출력하면 됩니다.

 

K가 12라고 가정하면,

arr = [1, 2, 2, 3, 3, 4, 4, 4, 5, 5, 6, 6, 8, 8, 9, 10, 10, 12, 12, 15, 15, 16, 20, 20, 25]

arr[12] // 6

 

위와 같이, 6을 출력하면 됩니다.

이분 탐색

정답은 이분 탐색입니다.

 

행렬을 직접 만든 다음 정렬을 시키려고 하면 1분 넘게 걸리니 포기합시다.

 

처음에는 도대체 이게 왜 이분 탐색인데? 라는 생각이 많이 들었는데, 내가 이분 탐색 문제를 안 풀어서 그런 것 같습니다.

 

문제를 푸는 데에는 몇 가지 키 포인트가 있습니다.

1. arr[K]의 값은 K가 2이상이고 N**2 미만일 때 무조건 K보다 작다. (arr[K] <= K, 단, 2 <= K < N**2)
2. arr[K]의 값보다 작은 값이 K-1개가 있다면, 그것이 정답이다. 따라서 arr[K]값을 이분 탐색으로 추측해가며 값의 범위를 좁혀나가자.

1번의 키포인트에 따라, arr[K]의 값의 한계는 '사실상' K이므로, 이분 탐색의 한계는 K로 설정하면 됩니다.

 

start를 1로, end를 K로 둔 뒤, 둘을 더하고 2로 나누어 추측을 시작할 중간값 mid를 만듭니다.

 

그렇다면, 어떻게 mid보다 작은 값을 구할 수 있을까요?

cnt = 0 # mid보다 작은 값의 수 카운트

for i in range(1, N+1):
    for j in range(1, N+1):
        if i*j <= mid:
            cnt += 1

 

위와 같은 방식으로 구할 수는 있으나, 연산량이 N^2가 되어버립니다.

 

행렬은 정사각형에 좌상단-우하단 대각선을 기준으로 대칭된다.

 

따라서, 첫 번째 열은 1로 시작하여 1씩 증가합니다.

 

두 번째 열은 2로 시작하여 2씩 증가합니다.

 

세 번째 열은 3으로 시작하여 3씩 증가합니다.

 

mid를 i(열의 번호)로 나누면, 보다 적은 연산량으로 mid보다 작은 값이 몇 개가 있는지 파악할 수 있습니다.

 

다만, mid가 N을 초과할 경우, 해당 열의 원소 개수를 초과하여 값을 더해버릴 수 있기 때문에, 더할 수 있는 최대 값은 N으로 설정합니다.

 

cnt = 0

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

 

이제, cnt가 K와 같거나 넘은 경우, 넘지 못한 경우를 따지면 됩니다.

1. cnt가 K와 같거나 넘었다?

같다면 답이 될 수 있지만, 이분 탐색이 끝나기 전에는 잘못된 값이 나올 수 있으므로 일단 answer에 mid를 할당합니다.

넘었다면, 설정한 추측값(mid)가 높으므로, end를 mid로 설정하여 낮추고, 1을 빼줍니다.

2. cnt가 K를 넘지 못했다?

설정한 추측값(mid)가 너무 낮으므로, start를 mid로 설정한 뒤, 1을 더합니다.

while문의 반복 조건으로 걸어두면 되는 것이죠.

 

cnt와 K의 값이 같다고 해도, start와 end의 범위가 넓을 경우,

조건은 만족하지만 정답이 아닌 값이 나올 수 있으므로 start가 end와 같거나 넘어설 때까지 반복해야 합니다.

정답 코드

N, K = int(input()), int(input())

s, e = 1, K
a = 0

while s <= e:

    m = (s + e) // 2
    cnt = 0
    for i in range(1, N+1):
        cnt += min(m//i, N)

    if cnt >= K:
        a = m
        e = m-1
    else:
        s = m+1

print(a)