1300번: K번째 수

1300번: K번째 수

사용 언어: C++

문제 요약

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

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

풀이

문제를 보고 처음 접근한 방식은 배열의 모양에서 규칙을 찾으려고 했습니다. i+j 값이 같은 대각선끼리 묶으면 뭔가 규칙이 있는 것 같아서 그런 식으로 가설을 세웠는데

"i+j 값이 작은 대각선에 있는 숫자들은 i+j 값이 큰 대각선의 숫자들보다 항상 작을 것이다", 를 가정하고 예시를 보는데, 보니까 값의 크기 순서는 대각선의 순서와 전혀 상관이 없더군요.

 

그래서 필요한 것이 발상의 전화!

배열을 직접 만들 수도 없는 상황이고, 위치에서 규칙을 찾을 수도 없다면, 탐색 같은 것을 하는 게 아닌, '정답이 될 수 있는 값'의 범위를 탐색해야 합니다. 그래서 이분 탐색이 여기서 사용됩니다.

 

k번째 수 B[k]는 어떤 값 X일 것인데, 이 X를 직접 찾는 건 어렵지만, X보다 작거나 같은 수가 배열에 총 몇 개가 있을까?,라는 질문은 답하기 쉽습니다. 이 질문을 통해 X가 정답보다 큰지 작은지를 판단할 수 있습니다.

예를 들어, 찾으려는 B[k]가 7번째 수라고 하면,

- 만약 X를 10으로 추측했는데, 10보다 작거나 같은 수가 5개뿐이라면? -> 7번째 수는 10보다 무조건 커야함.

- 만약 X를 20으로 추측, 20보다 작거나 같은 수가 9개라면 -> 7번째 수는 20보다 작거나 같아야함.

이렇게 계속 이분 탐색으로 범위를 줄여나가면 답을 시간 안에 찾을 수 있음.

 

해당 문제에서 이분 탐색의 핵심 논리

1. 탐색 범위 설정 : left=1, right = k (k번째 수는 k보다 클 수 없음)

2. mid = (left + right) / 2

3. 개수(count) 계산 : mid 보다 작거나 같은 수가 NxN 배열에 총 몇 개 있는지 계산. 여기서 이중 for문 같은걸 사용하는 게 아닌 효율적인 방법을 택해야하는데, 각 행 i에서 mid 보다 작거나 같은 수는 i*j <= mid를 만족해야 하므로, j <= mid / i. 따라서 각 행에는 min(N, mid/i)개의 숫자가 존재. 이걸 1행부터 N행까지 모두 더하면 count가 된다.

4. 범위 좁히기: count<k 이면 left = mid + 1로 설정. count >= k이면 ans = mid(임시 저장), right = mid-1;

5. 반복 및 종료

#include <iostream>
#include <algorithm>

using namespace std;

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    
    long long N, K;
    cin >> N >> K;
    
    long long left = 1;
    long long right = N * N;
    long long ans = 0;
    
    // 이분 탐색
    while(left <= right)
    {
        long long mid = left + (right - left) / 2;
        long long count = 0;
        
        // mid보다 작거나 같은 숫자의 개수
        for(int i = 1; i <= N; i++)
        {
            // i행 기준으로 열(j)의 최대 개수는 mid / i;
            // 근데 j는 N을 넘을 수 없어서 둘 중 최솟값.
            count += min((long long)N, mid / i);
        }
        
        if(count >= K)
        {
            // mid보다 작거나 같은 수가 K개 이상
            // mid가 답 or 답보다 큰 수. -> 일단 저장 후 다시 탐색
            ans = mid;
            right = mid - 1;
        }
        else
        {
            // mid보다 작거나 같은 수가 K개 미만
            // mid는 답보다 작은 수
            // 더 큰 값으로 재탐색
            left = mid + 1;
        }
    }
    
    cout << ans << '\n';

    return 0;
}