[C/C++] 백준 1300번 - K번째 수 (이분 탐색)

 

<코드>

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;

long long N, K;
long long Low, High, Mid;
long long cnt;

long long count(long long x)
{
	long long sum = 0;
	for (int i = 1; i <= N; i++)
	{
		sum += min(x / i, N);
	}
	return sum;
}

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

	K = min((long long)1000000000, K);

	Low = 1;
	High = N * N;

	while (Low <= High)
	{
		Mid = (Low + High) / 2;

		cnt = count(Mid); // Mid보다 작거나 같은 수들의 개수

		if (cnt >= K)
		{
			// Mid보다 작거나 같은 수들의 개수가 목표 인덱스보다 많으므로
			// 이므로 Mid값을 더 작게 찍어줘야 합니다.
			// 따라서 High값 조정
			High = Mid - 1;
		}
		else
		{
			// Mid보다 작거나 같은 수들의 개수가 목표 인덱스보다 작으므로
			// 이므로 Mid값을 더 크게 찍어줘야 합니다.
			// 따라서 Low값 조정
			Low = Mid + 1;
		}
	}
	cout << Low;
}

 

풀이 방법

단순히 문제만 놓고 본다면 A[N][N]의 배열을 B[N*N]배열에 옮겨놓고 sort함수로 정렬한 뒤 B [K]의 값을 출력하면 되는 문제라고 생각됩니다. 그러나 N의 입력이 100,000(십만)까지 가능하며 최악의 경우 N*N이 10^10(백억)이므로 무조건 시간 초과가 발생하게 됩니다. 그리고 A[100001][100001]의 배열 또한 메모리 초과로 사용할 수 없습니다. 따라서 단순히 배열을 저장하고 정렬을 하게 된다면 문제를 풀 수 없게 됩니다. 따라서 이분 탐색으로 빠르게 답을 구할 수 있는 알고리즘을 생각했습니다.

 

<알고리즘>

1. Low=1, High=N*N으로 설정

2. Low <= High 일 동안 while문 루프

   1. Mid값 구하기

   2. Mid값보다 작거나 같은 수들 구하기 => cnt

      1. cnt값이 목표 인덱스(K) 보다 크거나 같을 경우

         Mid를 좀 더 왼쪽으로 옮겨야 합니다.

         1. High = Mid - 1

      2. cnt값이 목표 인덱스(K) 보다 작을 경우

         Mid를 좀 더 오른쪽으로 옮겨야 합니다.

         1. Low = Mid + 1

3. 루프 탈출 후 Low값 출력

 

<문제 예시> 입력이 N = 4, K = 7일 때를 생각해봅시다.

A[4][4]	[1]	[2]	[3]	[4]
[1]	1	2	3	4
[2]	2	4	6	8
[3]	3	6	9	12
[4]	4	8	12	16

먼저 풀이를 진행하기에 앞서 Mid값을 구했을 때 어떻게 Mid보다 작은 수들의 개수를 시간 초과 없이 구할 수 있는지부터 설명드리겠습니다. 

만약 처음의 Low = 1, High = 4*4 = 16으로 주어졌을 때 Mid값은 8입니다, 그리고 코드에서 보시면 count함수를 따로 빼서 구현을 했습니다.

long long count(long long x)
{
	long long sum = 0;
	for (int i = 1; i <= N; i++)
	{
		sum += min(x / i, N);
	}
	return sum;
}

count(Mid)로 count함수에서는 x값이 Mid값이 됩니다.

먼저 i가 1일 때 A의 1행은 1,2,3,...... 이런 식으로 1의 간격으로 증가하며 x보다 작은 수들은 1,2,3,..... 7,8로
x / i =8개입니다. 그러나 A의 1행에서 나올 수 있는 개수의 최댓값은 N개 이므로 min(x/i , N) = 4

i가 2일 때 A의 2행은 2,4,6,...... 이런 식으로 2의 간격으로 증가하며 x보다 작은 수들은 2,4,6,8로 
x / i =4개입니다. A의 2행에서 나올 수 있는 개수의 최댓값은 N개 이므로 min(x/i , N) = 4

i가 3일 때 A의 3행은 3,6,9,...... 이런 식으로 3의 간격으로 증가하며 x보다 작은 수들은 3,6으로 
x / i =2개입니다. A의 3행에서 나올 수 있는 개수의 최댓값은 N개 이므로 min(x/i , N) = 2

i가 4일 때 A의 2행은 4,8,...... 이런 식으로 4의 간격으로 증가하며 x보다 작은 수들은 4,8로 
x / i =2개입니다. A의 4행에서 나올 수 있는 개수의 최댓값은 N개 이므로 min(x/i , N) = 2

따라서 Mid = x = 8보다 작거나 같은 수들의 개수는 = 4 + 4 + 2 + 2 = 12가 됩니다.
이런 식으로 구현한다면 시간 복잡도 O(N) 만에 원하는 값을 얻을 수 있습니다. 

확인해볼까요? 배열 B를 생성하고 정렬했을 때 다음과 같은데

8보다 작거나 같은 수들의 개수를 세보면 12개가 맞습니다.

 

이제 문제로 돌아와서 N = 4, K=7일 때 풀이 과정을 살펴보면,

Low	High	Mid	count(Mid)	Mid vs K
1	16	8	12	12 > 7

초기의 Low와 High은 1과 4*4로 지정해주고 Mid값은 8입니다. Mid보다 작거나 같은 수의 개수는 12개입니다.

그러나 찾고자 하는 K는 12보다 작은 7이므로 Mid값을 더 작게 설정할 필요가 있습니다.

따라서 High을 Mid - 1로 설정해줍니다.

 

Low	High	Mid	count(Mid)	Mid vs K
1	16	8	12	12 > 7
1	7	4	8	8 > 7

찾고자 하는 K는 8보다 작은 7이므로 Mid값을 다시 더 작게 설정할 필요가 있습니다.

따라서 다시 High을 Mid - 1로 설정해줍니다.

 

Low	High	Mid	count(Mid)	Mid vs K
1	16	8	12	12 > 7
1	7	4	8	8 > 7
1	3	2	3	3 < 7

이번에는 count(Mid) 값이 K보다 작게 나왔습니다. 따라서 Mid값을 높게 설정해야 합니다.

Low를 Mid + 1로 설정해줍시다.

 

Low	High	Mid	count(Mid)	Mid vs K
1	16	8	12	12 > 7
1	7	4	8	8 > 7
1	3	2	3	3 < 7
3	3	3	5	5 < 7

이번에도 count(Mid) 값이 K보다 작게 나왔습니다. 따라서 Mid값을 높게 설정해야 합니다.

Low를 Mid + 1로 설정해줍시다.

 

Low	High	Mid	count(Mid)	Mid vs K
1	16	8	12	12 > 7
1	7	4	8	8 > 7
1	3	2	3	3 < 7
3	3	3	5	5 < 7
4	3	3	5	5 < 7

이번에도 count(Mid) < K입니다. 그러나 다시 루프를 돌려고 while문의 맨 위로 올라가려는 순간 while(Low <=High)이라는 조건으로 인해 while문을 탈출하게 되고 Low값을 출력하면 정답이 되게 됩니다.

 

 

 

 

이상입니다. 감사합니다.

 

 

 

www.acmicpc.net/problem/1300

1300번: K번째 수