BOJ 11004 : K번째 수

알고리즘 문제 전체 링크

 

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

 

 

O(NlogN) 정렬로 정렬한 후, A[K]를 출력하면 된다. 여기서는 merge sort를 사용하였다.

이때, K는 입력을 받은 후, 1을 빼야한다.

#include <stdio.h>

int N, K;
int A[5050000];
int B[5050000];

void merge(int* A, int start, int end)
{
	int mid, i, j, k;

	mid = (start + end) >> 1;
	i = start;
	j = mid + 1;
	k = 0;

	while (i <= mid && j <= end)
	{
		if (A[i] <= A[j]) B[k++] = A[i++];
		else B[k++] = A[j++];
	}

	while (i <= mid) B[k++] = A[i++];
	while (j <= end) B[k++] = A[j++];

	for (i = start; i <= end;i++)
		A[i] = B[i - start];

}

void sort(int* A, int start, int end)
{
	int mid;
	if (start >= end) return;

	mid = (start + end) >> 1;
	sort(A, start, mid);
	sort(A, mid + 1, end);
	merge(A, start, end);
}

int main(void)
{
	scanf("%d %d", &N, &K);

	K -= 1;
	for (int i = 0; i < N; i++) scanf("%d", &A[i]);

	sort(A, 0, N - 1);

	printf("%d\n", A[K]);

	return 0;
}

---

이 문제는 c++의 algorithm에서 nth_element를 이용할 수도 있다.

nth_element는 부분적으로 정렬을 시킨다.

배열 A[0]에서 A[N]까지 K번째 수를 찾고 싶다면, nth_element(A, A + K, A + N); 를 호출하면 된다.

그러면 A[K]에 K번째 수가 담겨있게 된다.

#include <stdio.h>
#include <algorithm>

using namespace std;

int N, K;
int A[5050000];

int main(void) 
{
	scanf("%d %d", &N, &K);

	K -= 1;
	for (int i = 0; i < N; i++) scanf("%d", &A[i]);

	nth_element(A, A + K, A + N);

	printf("%d\n", A[K]);

	return 0;
}