백준 알고리즘] 11004번 : K번째 수

시간 제한 : 1.5 초

메모리 제한 : 256 MB

문제

수 N개 A1, A2, ..., AN이 주어진다. A를 오름차순 정렬했을 때, 앞에서부터 K번째 있는 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N(1 ≤ N ≤ 5,000,000)과 K (1 ≤ K ≤ N)이 주어진다.

둘째에는 A1, A2, ..., AN이 주어진다. (-109 ≤ Ai ≤ 109)

출력

A를 정렬했을 때, 앞에서부터 K번째 있는 수를 출력한다.

소스코드

#include <iostream>
using namespace std;
/* Swap */
void Swap(int *a, int *b)
{
	int tmp = *a;
	*a = *b;
	*b = tmp;
}

/* Partition */
int Partition(int arr[], int left, int right)
{
	// 기본적으로 QuickSort의 Pivot 선정 후 교환 과정과 동일
	// pivot은 가장 왼쪽의 인덱스로 결정
	int pos = arr[left];
	int low = left + 1, high = right;
	while (low <= high)
	{
		while (low <= right && arr[low] <= pos)
		{
			low++;
		}
		while (high >= (left + 1) && arr[high] >= pos)
		{
			high--;
		}
		if (low <= high)
		{
			Swap(&arr[low], &arr[high]);
		}
	}
	Swap(&arr[left], &arr[high]);
	return high;
}

/* QuickSelect */
int QuickSelect(int arr[], int left, int right, int select)
{
	// 찾고자 하는 인덱스가 범위 내에 존재할 때,
	if (select > 0 && select <= right - left + 1)
	{
		int pos = Partition(arr, left, right);

		// 찾고자 하는 숫자를 찾은 경우
		if (pos - left == select - 1)
		{
			return arr[pos];
		}
		// 찾고자 하는 숫자가 small group에 있는 경우
		else if (pos - left > select - 1)
		{
			return QuickSelect(arr, left, pos - 1, select);
		}
		// 찾고자 하는 숫자가 large group에 있는 경우
		else
		{
			return QuickSelect(arr, pos + 1, right, select - pos + left - 1);
		}
	}
	// 탐색 실패 시,
	return -1;
}

int main(void)
{
    std::ios_base::sync_with_stdio(false);
	int n, index; // n: arr size, k: index
	int* arr;

	cin >> n;
	cin >> index;

	arr = new int[n];

	for (int i = 0; i < n; i++)
	{
		cin>>arr[i];
	}

	cout << QuickSelect(arr, 0, n - 1, index);
	return 0;
}

문제 풀이 과정

고민할 것 없이 그냥 QuickSort를 응용한 알고리즘인 QuickSelect 알고리즘을 이용했다.

STL에서 제공하는 잘 최적화된 정렬을 이용하여 K번째 수를 선택하는 것도 버퍼 관련 함수를 적절히 사용하면, 복잡한 알고리즘을 사용하지 않아도 되지만, 그래도 알고리즘을 공부하는 입장이라면 처음 생각한 방법 외의 것을 고민해볼 필요가 있다.