[python] 백준 1300 - K번 째 수
백준 1300 - K번 째 수 [GOLD 1]
두 개의 입력이 들어온다. N과 K.
길이가 N인 정사각형 형태의 행렬이 arr이 있고, arr[i][j]에는 i*j가 존재한다.
단, 인덱스는 1부터 시작한다.
즉, N이 5일 때, 행렬은 다음과 같이 생성된다.
1
2
3
4
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라고 가정하면,
1
2
3
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분 넘게 걸리니 포기하자.
처음에는 도대체 이게 왜 이분 탐색인데? 라는 생각이 많이 들었는데, 내가 이분 탐색 문제를 안 풀어서 그런 것 같다.
문제를 푸는 데에는 몇 가지 키 포인트가 있다.
arr[K]의 값은 K가 2이상이고 N**2 미만일 때 무조건 K보다 작다. (arr[K] <= K, 단, 2 <= K < N**2)
arr[K]의 값보다 작은 값이 K-1개가 있다면, 그것이 정답이다. 따라서 arr[K]값을 이분 탐색으로 추측해가며 값의 범위를 좁혀나가면 된다.
1번의 키포인트에 따라, arr[K]의 값의 한계는 ‘사실상’ K이므로, 이분 탐색의 한계는 K로 설정하면 된다.
start를 1로, end를 K로 둔 뒤, 둘을 더하고 2로 나누어 추측을 시작할 중간값 mid를 만든다.
그렇다면, 어떻게 mid보다 작은 값을 구할 수 있을까?
1
2
3
4
5
6
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으로 설정한다.
1
2
3
4
cnt = 0
for i in range(1, N+1):
cnt += min(mid/i, N)
이제, cnt가 K와 같거나 넘은 경우, 넘지 못한 경우를 따지면 된다.
- cnt가 K와 같거나 넘었다?
같다면 답이 될 수 있지만, 이분 탐색이 끝나기 전에는 잘못된 값이 나올 수 있으므로 일단 answer에 mid를 선언한다.
넘었다면, 설정한 추측값(mid)가 높으므로, end를 mid로 설정하여 낮추고, 1을 빼준다.
- cnt가 K를 넘지 못했다?
설정한 추측값(mid)가 너무 낮으므로, start를 mid로 설정한 뒤, 1을 더한다.
while문의 반복 조건으로 걸어두면 되는 것이다.
cnt와 K의 값이 같다고 해도, start와 end의 범위가 넓을 경우,
조건은 만족하지만 정답이 아닌 값이 나올 수 있으므로
start가 end와 같거나 넘어설 때까지 반복해야 한다.
정답 코드
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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)