[백준/C++] 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(10^9, N^2)\)보다 작거나 같은 자연수이다.

출력

B[k]를 출력한다.

---

접근법

많이 부족해서 정말 많이 고민하고 틀리고, 애를 썼던 문제이다.

문제를 먼저 설명하자면, N x N크기의 배열 A의 원소를 1차원 배열에 오름차순 정렬했을 때, k번째 오는 값을 구하는 것이다. 예를 들어 N = 3, k = 7에서 원소를 1차원 정렬을 진행하면 1 2 2 3 3 4 6 6 9 이므로 7번째 값인 6이 답이 된다.

해당 문제는 모든 경우의 수를 탐색하는 브루트 포스를 사용하면 정말 쉽게 풀 수 있다. 하지만, N의 최댓값이 \(10^5\)이며, 이를 가정했을 때 k의 값도 \(10^9\)이므로 메모리 제한을 넘기 때문에 전체를 순회하는 방식은 사용할 수 없다.

브루트 포스가 아닌, 시간복잡도를 크게 줄일 수 있는 이진탐색을 사용할 수 있다. 이진탐색 적용 조건은 탐색하려는 배열이 정렬되어있어야하는 것이다. 정렬되어있지 않으면 원하는 값이 나눌 두 개의 부분 중 어디에 있는지 보장할 수 없기 때문이다.

이진탐색을 사용할 수 있는 근거는 다음과 같다. 이진탐색을 수행할 배열의 정의를 ‘x숫자보다 작거나 같은 수의 개수’라고 하는 것이다.

N = 3일때 A[3][3] = {1, 2, 2, 3, 3, 4, 6, 6, 9}이고, 배열의 이름을 arr라고 하겠다. 배열은 다음과 같다.

arr[1] = 1 arr[2] = 3 arr[3] = 5 arr[4] = 6 arr[5] = 6 arr[6] = 8 arr[7] = 8 arr[8] = 8 arr[9] = 9

다음과 같이 인덱스의 숫자가 늘어날수록 배열값들은 우상향함을 확인할 수 있다. 인덱스가 1씩 증가해도 배열값이 같은 경우도 존재한다.

k = 7이라고 하면 A배열에서 6이 2개 들어있어 8을 배열값으로 가진 arr[6], arr[7]이 존재하기 때문에 구하려는 값이 6이 나올수도 있고 7이 나올수도 있다. 이때 주목해야 할 가장 중요한 포인트는 다음과 같다.

arr[6] = 8, arr[7] = 8의 의미는 6보다 작거나 같은 수의 개수는 여덟 개, 7보다 작거나 같은 수의 개수는 여덟 개이다. 즉, 이 말은 7이 아닌 6이 한번 더 A배열에 존재함을 뜻한다.

7이 A배열에 없다는건 해당 숫자를 고려할 필요가 없다는 것이다. 그러므로 k = 7 번째 수를 찾는다고 하였을때 우리는 arr[x] == 7를 만족하는 가장 작은 수를 찾아야함을 알 수 있다.

이를 구현한 코드는 다음과 같다.

---

코드

#include <iostream>
using namespace std;

long long N, k;

int main() {
    cin >> N;
    cin >> k;

    long long start = 1;
    long long end = N * N;
    long long mid, answer, cnt;

	while (start <= end) {
        cnt = 0;
		mid = (start + end) / 2;

		for (long long i = 1; i <= N; i++) {
			cnt += min(N, mid / i);
		}
		
		if (cnt < k) start = mid + 1;
		else {
            answer = mid;
            end = mid - 1;
        }
	}
    cout << answer;
}

cnt < k인 경우에는 확실히 오른쪽(end 부분)에 답이 있음을 보장하기 때문에 start = end + 1로 구현하였다.

반대로 cnt >= k인 경우에는 위에서 언급했듯이 해당 위치(mid)가 답일 수도 있고, 값이 중복되어 있어 왼쪽 부분에도 있을 수 있으므로 일단 answer = mid로 값을 저장해놓고 end = mid - 1;으로 왼쪽 부분으로 이동한다.

while문 조건이 start <= end 이므로 해당 조건을 반할때까지 반복하면 cnt == k를 만족하되 인덱스값이 가장 작은 값을 return할 수 있다.

---

느낀점

이진탐색을 너무 쉽게 봤던 것 같다. 분명 쉬운 알고리즘이지만 이를 도출해내기까지 오랜 시간이 걸렸다. 주어진 조건을 어떻게 해석하냐가 중요하다고 생각한다.

---

원문: Velog