K번째 수

백준 11004번 K번째 수

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

입력

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

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

출력

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

예제 입력 1	예제 출력 1
5 2	2
4 1 2 3 5

풀이

pivot을 정하는 방법

• pivot == K： K번째 수를 찾은 것이므로 알고리즘을 종료한다.
• pivot > K： pivot의 왼쪽 부분에 K가 있으므로 왼쪽( ~ pivot - 1)만 정렬을 수행한다.
• pivot < K： pivot의 오른쪽 부분에 K가 있으므로 오른쪽(pivot - 1 ~ )만 정렬을 수행한다.

소스

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		final int N = Integer.parseInt(st.nextToken()); //숫자 개수
		final int K = Integer.parseInt(st.nextToken()); //K번째 수
		st = new StringTokenizer(br.readLine());
		int input [] = new int[N]; //입력받은 배열
		for(int i = 0; i < N; i++)
			input[i] = Integer.parseInt(st.nextToken());
		quickSort(input, 0, N - 1, K - 1);
		System.out.print(input[K - 1]);
	}

	static void quickSort(int[] nums, int start, int end, int k) {
		if(start < end) {
			int pivot = getPivot(nums, start, end);
			if(pivot == k)
				return;
			else if(pivot > k)
				quickSort(nums, start, pivot - 1, k);
			else
				quickSort(nums, pivot + 1, end, k);
		}
	}

	static int getPivot(int[] nums, int start, int end) {
		if(start + 1 == end) {
			if(nums[start] > nums[end])
				swap(nums, start, end);
			return end;
		}
		int middle = (end + start) / 2;
		
		swap(nums, start, middle);
		
		int pivot = nums[start];
		int i = start + 1;
		int j = end;
		while(i <= j) {
			while(nums[j] > pivot && j > 0) 
				j--;
			while(nums[i] < pivot && i < nums.length - 1) 
				i++;
			if(i <= j) {
				swap(nums, i++, j--);
			}
		}
		nums[start] = nums[j];
		nums[j] = pivot;
		return j;
	}

	static void swap(int[] nums, int i, int j) {
		int tmp = nums[i];
		nums[i] = nums[j];
		nums[j] = tmp;
	}
	
}