k번째 수

한동안 코딩테스트 관련한 포스팅을 게을리하다, 이 문제는 다시 한 번 풀어도 헤맬 것 같아 미래의 나를 위해 정리합니다.

🛠문제

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

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

N = 3, K = 7
A = [ [1, 2, 3], [2, 4, 6], [3, 6, 9] ]
B[7]에 해당하는 수는?

💻소스코드

N, K = int(input()), int(input())
sp, ep = 1, K

while sp <= ep:
    md = (sp + ep) // 2

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

    if cnt >= K:
        answer = md
        ep = md - 1
    else:
        sp = md + 1

print(answer)

✏️풀이 설명

이분 탐색 문제인 것을 알면서도 왜 이분 탐색인지 알지 못했던 문제입니다... 너무 어렵다....

이 문제를 풀 때 가장 이해하기 어려웠던 부분은 다음 두 가지입니다.

1. k번째 수는 k를 넘을 수 없습니다.
2. A는 i의 배수로 이루어진 배열입니다. 어떤 수 m 이하의 수의 개수는 m / i의 몫으로 구할 수 있습니다.

1번의 이유는 1부터 k까지 모든 수는 적어도 한 번이상 등장합니다. A의 배열은 배수로 이루어진 배열이기 때문입니다. 이 부분이 의심스럽다면 end 값을 최대값으로 설정해도 되긴합니다.

2번은 다음과 같이 생각할 수 있습니다.

만약, A 배열에서 7이하인 수가 몇개인지 알고싶다고 가정합시다. A의 배열은 i의 배수로 이루어진 배열이고 A 배열 요소를 i로 나눈다면 모두 동일하게 [1, 2, 3]입니다. (N이 3일때)

i가 2일때 A 배열은 [2, 4, 6]입니다. 이 중에서 7이하의 수를 어떻게 구할까요? 이 배열을 i로 나누면 [1, 2, 3]이고 7을 2로 나누면 3.5입니다. 배열은 모두 정수이므로 소수점을 버리면 7 이하의 수는 3개입니다.

코드에서 min(mid // i, N) 이라는 표현이 나왔는데, 이는 배열의 크기가 N개인데, 이 배열에서 7보다 작은 수의 개수가 N 이상이 나올 수 없기 때문입니다.

i가 1일때 7을 1로 나누면 몫은 7입니다. 이 말은 7 이하의 수가 7이라는 말인데, 배열의 크기가 3이므로 7이하의 수는 최대 3개만 가능하므로 min()함수를 사용했습니다.

단순히, 각 배열 요소가 i의 배수이니까 i로 나누었고, 비교해주기 위해 7도 i로 나누어 비교한 것 뿐입니다.

a = b
a / x = b / x