[백준 11004] K번째 수 - Quick select

www.acmicpc.net/problem/11004

11004번: K번째 수

 

 

단순히 정렬해서 풀 수가 없었기 때문에, 이것저것 찾아보다가 Quick select라는 알고리즘을 발견하게 됐습니다.

뭐 그리 대단한건 아니고 퀵 소트에 사용된 연산을 약간 응용한 버전입니다.

 

퀵 소트에서 피봇이 정해지면 피봇의 값을 기준으로 왼쪽에는 그보다 작은 값, 오른쪽에는 그보다 큰 값이 오도록 배열의 값을 수정하는 연산이 있습니다.  자세한 건 퀵 소트 관련글들을 참고해주세용. (링크 - 퀵 소트 참고글)

 

중요한 점은 한번의 연산이 끝날 때마다 피봇의 위치가 정해진다는 겁니다. 

 

잠시 간단하게 생각해볼까요?

피봇을 정하고 그보다 작은 값을 모두 왼쪽으로 보내고 큰 값은 모두 그보다 오른쪽으로 보내 정렬한다고 했을때, 전체 배열이 정렬되지는 않았더라도 피봇이 왼쪽의 요소들보다 크므로 n번째(왼쪽 요소 배열의 길이가 n-1일 때) 크다고 할 수 있습니다. 그리고 이 자리는 전체배열에서 그대로 확정되는 값입니다.

 

그러면 한 번의 연산이 끝나고 피봇의 인덱스가 k보다 작다면 오른쪽 배열을 선택하고, 더 크다면 왼쪽 배열을 선택합니다.

이진 탐색의 경우와 비슷하다고 생각하면 됩니다. 이진 탐색의 경우에는 정렬된 배열에서만 수행가능 했는데, 이는 정렬되지 않는 배열에서도 사용할 수 있다는 점이 다른 점이네요. 

 

Quick select의 경우 최악의 케이스에서는 O(n^2)의 복잡도를 보여줍니다. 

출처 : https://www.wikiwand.com/en/Quickselect

주어진 배열에서 피봇을 선택할 때마다 최대값 혹은 최소값만 선택하는 경우엔  O(N^2)이 될 수 있습니다. 

제 경우에는 최악의 케이스를 방지하기 위해 피봇 선정시 난수를 이용해 랜덤하게 선정되도록 했습니다.

 

---

입력시 StrignTokenizer와 String.split 의 성능 차이에 대해 정리도 해둬야겠습니다. 

 

 

결론은 StringTokenizer가 훨~씬 빨랐습니다. 1초대는 StrignTokenizer를 사용했을 때고, 4초대는 String.split을 사용했을 때입니다. 두 경우 모두 BuffredReader를 사용해서 입력을 받았습니다.

 

 

import java.io.*;
import java.util.Random;
import java.util.StringTokenizer;

/*
백준 11004번 - K번째 수
 */
public class Main {
	static int n,k;
	static Random random = new Random();
	public static void main(String[] args) throws IOException {
//		InputReader reader = new InputReader("testcase.txt");
		InputReader reader = new InputReader();
		StringTokenizer st = new StringTokenizer(reader.readLine());

		n = Integer.parseInt(st.nextToken());
		k = Integer.parseInt(st.nextToken())-1;
		int[] a = new int[n];

		st = new StringTokenizer(reader.readLine());
		for (int i = 0; i < n; i++) {
			a[i] = Integer.parseInt(st.nextToken());
		}
		System.out.println(quickSelection(a,k));
//		quickSelection(a, 0 ,n-1);
//		System.out.println(a[k]);
	}
//  재귀 ver
//	private static void quickSelection(int[] a, int start, int end) {
//		if(start == end) return;
//
//		int pivotIdx = getPivotIndex(a,start,end);
//
//		if(pivotIdx < k) quickSelection(a, pivotIdx+1, end);
//		if(pivotIdx > k) quickSelection(a,start,pivotIdx-1);
//		if(pivotIdx == k) return;
//	}

	private static int quickSelection(int[] a, int k) {
		int start = 0; int end = a.length-1;
		while(true) {
			int pivotIdx = getPivotIndex(a, start, end);
			if(pivotIdx < k){
				start = pivotIdx+1; continue;
			}
			if(pivotIdx > k) {
				end = pivotIdx-1; continue;
			}

			if(start == end) return a[start];
			if(pivotIdx == k) {
				return a[pivotIdx];
			}
		}
	}

	private static int getPivotIndex(int[] a,int start, int end) {
		swap(a,start, getRandomInt(start, end));

		int pivotIdx = start;
		int l = start; int h = end;

		while(l < h) {
			while(a[h] > a[pivotIdx]) {
				h--;
			}

			while(l < h && a[l] <= a[pivotIdx]) {
				l++;
			}
			swap(a,l,h);
		}
		swap(a,pivotIdx, l);
		return l;
	}

	private static int getRandomInt(int start, int end) {
		if(start == end) return start;
		int offset = end-start;
		return random.nextInt(offset)+start;
	}

	private static void swap(int[] a,int n1, int n2) {
	    int temp = a[n1];
	    a[n1] = a[n2];
	    a[n2] = temp;
	}
}

class InputReader {
	private BufferedReader br;

	public InputReader() {
		br = new BufferedReader(new InputStreamReader(System.in));
	}

	public InputReader(String filepath) {
		try {
			br = new BufferedReader(new FileReader(filepath));
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		}
	}

	public void close() throws IOException {
		br.close();
	}

	public String readLine() throws IOException {
		return br.readLine();
	}
	public int readInt() throws IOException {
		return Integer.parseInt(readLine());
	}
	public Long readLong() throws IOException {
		return Long.parseLong(readLine());
	}
}