[백준/Java] 1300 K번째 수

https://www.acmicpc.net/problem/1300

1. 문제 요약

크기가 N x N인 배열 A에서 A[i][j] = i * j이다.
이 값을 모두 1차원 배열 B에 넣은 후 오름차순 정렬을 했을 때, k번째 수인 B[k]를 구하여라.
단, A와 B의 인덱스는 1부터 시작한다.

배열의 크기 N은 최대 10만이다.
찾고자 하는 숫자의 인덱스 k는 최대 10억이다.

2. 문제 접근

가장 먼저 고민한 점은 10억번째 숫자의 값이 20억을 넘어서 int형을 사용하지 못하는 경우가 있는 지였다.
그런데 $n^2$의 형태일 때가 가장 가능한 큰 값을 가질텐데, 그런 경우라면 그냥 $n^2$번째 값이 $n^2$의 값을 갖기 때문에 그 이하이고, int형을 사용할 수 있다는 결론을 얻는다.

다만 이번에는 시간복잡도가 문제가 된다.
배열의 크기 N이 최대 10만이기 때문에 N x N 크기의 배열 A를 모두 구현하는 시간복잡도 $O(n^2)$를 감당할 수 없다. 즉, 배열 A와 B를 직접적으로 구현할 수 없다.
그러면 N과 k만으로 값을 구할 수 있어야 하는데, 어떻게 구할 수 있을까?

해답은 이분 탐색이었다.
앞서 이야기했듯이 k번째 수는 k이하의 값을 갖는다.
따라서 숫자 1부터 k까지 이분탐색을 하면서 각 값이 위치하는 인덱스를 판단하는 방식으로 접근할 수 있다.

다만 해당 값이 한 개보다 많이 있을 수 있다거나 하나도 없는 경우가 문제가 된다.
하지만 해당 값보다 작은 값을 찾는다는 접근으로 생각하면 이러한 문제도 해결할 수 있다.
배열 A를 세로로 훑으며 해당 값 이하의 개수, 그리고 (해당 값-1) 이하의 개수를 찾아서 해결할 수 있다.
배열 A의 값이 두 인덱스의 곱으로 얻어지기 때문에 각 row에서 해당 값 이하를 얻는 것은 단순한 나머지 연산($O(1)$)으로 수행할 수 있다.
이 부분은 1부터 N까지 탐색하게 된다.

최종적으로 이분 탐색을 수행하는 시간복잡도 $O(logk)$와 탐색 값 이하의 개수를 찾는 시간복잡도 $O(N)$을 합친 전체 시간복잡도는 $O(Nlogk)$이므로 문제를 해결할 수 있다.

3. 풀이

기본적인 아이디어는 다음과 같다.

1. 1부터 k까지의 이분탐색을 수행한다.
2. 이때 현재 탐색하는 값의 시작 인덱스와 끝 인덱스를 얻기 위해 먼저 탐색 값 이하의 개수, 그리고 (탐색 값-1) 이하의 개수를 찾는다.
3. K가 탐색하는 값의 끝 인덱스보다 크면 오른쪽으로, 시작 인덱스보다 작다면 왼쪽으로 넘어간다.

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

import java.util.*;
import java.io.*;

public class Main {

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        int N = Integer.parseInt(br.readLine());
        int K = Integer.parseInt(br.readLine());

        int left = 1;
        int right = K;

        // B[K]는 무조건 K이하의 값을 갖기 때문에 1부터 K까지의 이분탐색을 수행한다.
        while (left < right) {
            int mid = (left + right) / 2;

            // mid 값의 인덱스 범위를 얻어 이동 방향을 결정한다.
            int midLastIndex = 0;
            int midFirstIndex = 0;

            for (int i = 1; i <= N; i++) {
                midLastIndex += Math.min(N, mid/i);
                midFirstIndex += Math.min(N, (mid-1)/i);
            }
            // 이전까지는 mid-1의 마지막 index를 나타내기 때문에 1을 더해줌
            midFirstIndex++;

            // midLastIndex보다 크다면 mid는 오른쪽으로 가야함
            if (K > midLastIndex) {
                left = mid+1;
            }
            // midFirstIndex보다 작다면 mid는 왼쪽으로 가야함
            else if (K < midFirstIndex) {
                right = mid-1;
            }
            // 만약 그 사이라면 이게 답
            else {
                left = mid;
                break;
            }
        }

        int answer = left;

        System.out.println(answer);
    }
}

4. 회고

여기서는 해당없지만, 일반적인 이분 탐색은 인덱스 처리가 항상 고민된다. 이분 탐색을 완벽히 공부해서 글을 하나 써야겠다.
틀렸던 이유: mid가 바로 답인 경우를 고려하지 않았다.