[백준] 1300번 K번째 수 (C++)

문제 링크

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

1300번: K번째 수

 

문제 설명

N과 K가 주어진다.

주어진 N에 대하여 NxN인 배열 A를 가정할 때, 배열에 들어있는 수 A[i][j] = ixj 이다.

이를 일차원 배열 B에 넣으면 NxN 크기가 되며, B를 오름차순 정렬 할 때 B[k]를 구한다.

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

 

사고 과정

가장 단순하게 생각한다면 NxN의 배열을 반복문을 통해 생성한 후,

이를 오름차순 정렬 후에 인덱스 k에 대해 접근하면 될 것이다.

하지만 이 방법의 경우 NxN의 배열을 생성하며 이미 O(N^2)이 되기 때문에 시간초과가 날 것이다.

따라서 배열을 생성하면 안되고, 특정 규칙에 따라 오름차순으로 정렬된 배열의 인덱스 접근을 해야 함을 알 수 있다.

 

해결 과정

B[k]에 해당하는 값을 S라고 할 때, 우리는 주어진 k에 해당하는 S를 찾아야 한다.

하지만 해당 문제를 해결하기 위해서는 반대로 생각해야 할 필요가 있다.

k 값을 이분 탐색을 활용하여 옳바른 S를 구하는 것이 아닌,

S를 이분 탐색으로 활용하여, 즉, S의 값을 조절하며 해당 S가 주어진 k와 동일한지 판단해야 한다.

 

백준에서 주어진 예시를 확장하여 생각해보자.

N이 4이고 S가 9라면 K는 어떤 값을 갖게 될까?

N이 4이고 S가 12라면 K는 이전과 어떻게 변할까?

4x4 배열에 대해, S가 9라면 각 열에 대해 자신보다 작거나 같은 수의 갯수는

각각 4, 4, 3, 2가 되므로, k는 13이 될 것이다.

S가 12라면, 각각 4, 4, 4, 3이 되므로 k는 15가 될 것이다.

 

따라서 우리는 S의 값을 1부터 N^2까지의 범위로 설정하고 이분 탐색을 통해 

올바른 k를 가지는 S의 값을 구하면 되며, S의 이분 탐색 시간은 O(logN)이 된다.

 

다만 시간 초과를 해결하기 위해서는 S가 주어졌을 때

NxN 배열에서 S 이하인 수의 개수를 찾는 과정이 O(N)의 시간 복잡도 안에서 해결되어야 한다.

위의 그림을 보면 알 수 있듯이, NxN 배열을 N개의 열에 대해 나누고,

각 열에서 S 이하인 숫자를 찾는 과정을 O(1)에 해당하는 알고리즘을 활용하면 된다.

 

이 문제의 핵심이라고도 할 수 있는데, 각 열에서 S 이하인 숫자를 찾기 위해서는 min(N, S/i)를 활용하면 된다. (i = 인덱스)

우선 각 열에서 최대로 가질 수 있는 S 이하인 숫자의 갯수는 N일 것이므로 (NxN 배열이므로), min(N, )이 필요하다.

또한 각 배열의 수가 ixj라는 점을 고려해보면, S를 해당 열의 인덱스로 나눈 값은 해당 열의 행 번호가 될 것이므로

각 행에서 몇 번째에 위치하는지를 알 수 있게 된다. (S/i 가 N보다 작을 경우에만)

이는 각 열에서 S가 몇 번째로 큰 지를 의미하므로, 단순 나눗셈을 통해 S가 각 열에서 몇 번째로 큰 숫자인지 알 수 있는 것이다.

 

다만 여기서 주의할 점은 위의 예시에서 S가 12일 때 k가 15라는 값이 도출되지만, 실제로는 k가 14일 때에도 S는 12의 값을 갖는다.

따라서 이분 탐색을 통해 S의 크기를 조절할 때, 주어진 S에 대한 k의 값 (코드에서 cnt 변수)이 입력 받은 k의 값보다 크거나 같다면

우리가 출력하고자 하는 값 (코드에서는 result 변수)를 바꿔주어야 한다.

만약 주어진 S에 대한 k의 값과 입력 받은 k의 값이 같을 때에만 출력하고자 하는 값을 변경하게 된다면

배열에 존재하지 않는 값이 출력되거나 무한 루프를 반복할 것이다.

 

코드

#include <iostream>
#include <algorithm>
#include <cmath>

using namespace std;

int N, K;

void init();
void input();
int count(int num);
void solve();

void init()
{
    ios::sync_with_stdio(false);
    cin.tie(NULL);
}

void input()
{
    cin >> N >> K;
}

int count(int num)
{
    int cnt = 0;
    for (int i = 1; i <= N; i++)
    {
        cnt += min(N, num / i);
    }
    return cnt;
}

void solve()
{
    int start = 1;
    int end = min(pow(10, 9), pow(N, 2));
    int mid, cnt, result;

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

        if (cnt < K)
        {
            start = mid + 1;
        }
        else if (cnt > K)
        {
            end = mid - 1;
            result = mid;
        }
    }
    cout << result;
}

int main()
{
    init();
    input();
    solve();
    return 0;
}