[백준] 15965번 K번째 소수: 에라토스테네스의 체 최적화 해보기 - C++

들어가며

 알고리즘 문제해결을 하다 보면 소수와 관련한 문제를 자주 만나볼 수 있습니다. 오늘은 소수판별법 중 하나인 에라토스테네스의 체를 몇 가지 방법으로 최적화해보겠습니다.

K번째 소수

 자연수 K가 주어지면 K번째 소수를 출력하는 간단한 문제이지만 K의 범위가 최대 50만이기 때문에 시간초과가 나오기 쉬운 문제입니다.

풀이방법

 기본적인 에라토스테네스의 체 알고리즘은 다음과 같습니다. 소수의 배수들을 제거해 나가면서 소수의 후보를 줄이는 방식입니다.

#include <iostream>
#include <vector>

using namespace std;

int main()
{
	vector<int> prime(500001); // 소수를 담는 배열. 처음엔 비어있음.
	int prime_num = 1; // 첫번째 소수를 입력할 차례.

	int limit = 7368788; // 50만번째 소수. 문제에서 가능한 가장 큰 소수
	vector<bool> isPrime(limit + 1, true); // 소수를 판별하는 배열
	isPrime[0] = isPrime[1] = false; // 0과 1은 소수가 아님

	for (int i = 2; i <= limit; i++)
	{
		if (isPrime[i]) // i가 소수라면 
		{
			prime[prime_num++] = i; // prime_num번째 소수는 i.
			for (int j = i * 2; j <= limit; j += i)
				isPrime[j] = false; // i의 배수는 소수가 아니다.
		}
	}

	int n;
	cin >> n;

	cout << prime[n]; // n번째 소수를 출력한다.

	return 0;
}

 

최적화 방법

 첫 번째 최적화 방법은 짝수를 제거하는 것입니다. 짝수인 소수는 2밖에 없기 때문에 검사할 때 짝수를 건너뛴다면 연산시간이 반으로 줄어듭니다.

#include <iostream>
#include <vector>

using namespace std;

int main()
{
	vector<int> prime(500001);
	prime[1] = 2; // 1번째 소수는 2
	int prime_num = 2; // 2번째 소수를 받을 차례

	int limit = 7368788;
	vector<bool> isPrime(limit + 1, true);
	isPrime[0] = isPrime[1] = false;

	for (int i = 3; i <= limit; i += 2) // 3부터 홀수만 검사한다.
	{
		if (isPrime[i])
		{
			prime[prime_num++] = i;
			for (int j = i * 2; j <= limit; j += i)
				isPrime[j] = false;
		}
	}

	int n;
	cin >> n;

	cout << prime[n];

	return 0;
}

 

 두 번째 최적화 방법은 내부루프를 i * i부터 도는 겁니다. 2부터 i - 1까지의 배수들을 앞서 제거했다는 사실을 생각해 보면 i의 배수중에 i * (i - 1) 까지는 이미 모두 제거되었다는 것을 알 수 있습니다. 

#include <iostream>
#include <vector>

using namespace std;

int main()
{
	vector<int> prime(500001);
	prime[1] = 2;
	int prime_num = 2;

	int limit = 7368788;
	vector<bool> isPrime(limit + 1, true);
	isPrime[0] = isPrime[1] = false;

	for (int i = 3; i <= limit; i += 2)
	{
		if (isPrime[i])
		{
			prime[prime_num++] = i;
			for (long long j = 1LL * i * i; j <= limit; j += i) // i * i 부터 제거한다. 
				isPrime[j] = false; // 이 과정에서 오버플로우가 날 수 있기에 j의 타입을 long long 으로 바꾼다.
		}
	}

	int n;
	cin >> n;

	cout << prime[n];

	return 0;
}

 

 세 번째 최적화 방법은 내부 루프 또한 홀수인 합성수만 제거하는 것입니다. 어차피 루프는 홀수만 돌기 때문에 홀수인 합성수만 제거해도 괜찮습니다.

#include <iostream>
#include <vector>

using namespace std;

int main()
{
	vector<int> prime(500001);
	prime[1] = 2;
	int prime_num = 2;

	int limit = 7368788;
	vector<bool> isPrime(limit + 1, true);
	isPrime[0] = isPrime[1] = false;

	for (int i = 3; i <= limit; i += 2)
	{
		if (isPrime[i])
		{
			prime[prime_num++] = i;
			for (long long j = 1LL * i * i; j <= limit; j += 2 * i) // 홀수 + 짝수 = 홀수 를 이용하여 2 * i씩 키워준다.
				isPrime[j] = false;
		}
	}

	int n;
	cin >> n;

	cout << prime[n];

	return 0;
}

 

 마지막 방법은 C++의 특성 때문인데 bool 배열 대신에 메모리 접근이 빠른 char배열을 사용하는 것입니다.

#include <iostream>
#include <vector>

using namespace std;

int main()
{
	vector<int> prime(500001);
	prime[1] = 2;
	int prime_num = 2;

	int limit = 7368788;
	vector<char> isPrime(limit + 1, true); // char배열로 변경
	isPrime[0] = isPrime[1] = false;

	for (int i = 3; i <= limit; i += 2)
	{
		if (isPrime[i])
		{
			prime[prime_num++] = i;
			for (long long j = 1LL * i * i; j <= limit; j += 2 * i)
				isPrime[j] = false;
		}
	}

	int n;
	cin >> n;

	cout << prime[n];

	return 0;
}

 

마치며

 오늘은 소수를 판별하는 에라토스테네스의 체를 다양한 방법으로 최적화해 보았습니다. 또 다른 최적화 방법이 있다면 댓글로 알려주시면 감사하겠습니다.