[백준 - 11004] K번째 수 .java

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

 

문제

수 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 

5 2
4 1 2 3 5

예제 출력 1 

2

 

Solution

n의 최대 범위( N ≤ 5,000,000 )가 크기 때문에 O(nlogn)의 시간 복잡도로 정렬을 수행하면 된다.

퀵 정렬을 이용해 구현하고자 했는데 너무 어려웠다.... 인터넷의 도움을 받았고 이해하는 것에 일단 만족했다.

 

quick 메서드

• quick 메서드는 배열의 특정 부분을 정렬하면서 k번째 작은 요소가 위치한 부분을 찾기
• partition 메서드를 호출하여 피벗을 기준으로 배열을 분할
• 피벗의 위치가 k와 같다면 해당 요소를 찾았으므로 메서드를 종료
• 그렇지 않으면, 피벗의 위치에 따라 배열의 왼쪽 또는 오른쪽 부분에서 계속 검색

partition 메서드

 

• partition 메서드는 주어진 배열 부분에서 피벗을 기준으로 작은 값과 큰 값을 분할
• 배열의 중앙값을 피벗으로 선택하고, 피벗과 첫 번째 요소를 교환
• 피벗을 기준으로 배열의 왼쪽과 오른쪽에서 작은 값과 큰 값을 찾고, 교환
• 마지막으로 피벗을 자신의 위치로 이동시키고 피벗의 위치를 반환

 

 

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());
        //변수 저장
        int n = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());

        //입력 받은 배열 저장
        int[] nums = new int[n];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < n; i++) {
            nums[i] = Integer.parseInt(st.nextToken());
        }

        quick(nums,0,n-1,k-1);
        System.out.println(nums[k-1]);
    }
    //배열 nums의 s부터 e까지의 부분을 정렬하며, k번째 작은 요소가 위치하는 부분을 찾기
    public static void quick(int[] nums, int s, int e, int k) {
        if(s<=e){
            int pivot = partition(nums,s,e);
            if(pivot == k){ // 더이상 정렬할 필요 없음
                return;
            } else if(pivot > k){ //pivot의 오른쪽만 다시 정렬
                quick(nums,s,pivot-1,k);
            } else{ //pivot의 왼쪽만 다시 정렬
                quick(nums,pivot+1,e,k);
            }
        }
    }
    //배열 nums의 s부터 e까지의 부분에서 피벗을 기준으로 작은 값과 큰 값을 분할
    public static int partition(int[] nums, int s, int e){
        if(s+1 == e){
            //배열의 크기가 2인 경우 처리
            if(nums[s] > nums[e]) {//직접 비교하여 정렬
                swap(nums,s,e);
            }
            return e;
        }
        // 중앙값을 피벗으로 선택
        int mid = (s + e) / 2;
        swap(nums, s, mid); // 피벗과 첫 번째 요소를 교환
        int pivot = nums[s]; // 피벗을 첫 번째 요소로 설정

        int i = s + 1; // 두 번째 요소부터 시작
        int j = e; // 마지막 요소부터 시작

        //i와 j 인덱스가 서로 교차하지 않는 한 계속해서 배열의 요소를 비교하고 교환
        while(i <= j){
            // 피벗보다 작은 값을 찾을 때까지 왼쪽으로 이동
            while( j>= s +1 && pivot < nums[j]) {
                j--;
            }
            // 피벗보다 큰 값을 찾을 때까지 오른쪽으로 이동
            while( i <= e && pivot > nums[i]) {
                i++;
            }
            // 값들이 엇갈리지 않은 경우 교환
            if( i <= j){
                swap(nums,i++,j--);
            }

        }
        // 피벗을 자신의 위치로 이동
        nums[s] = nums[j];
        nums[j] = pivot;
        return j;
    }
    public static void swap(int[] nums,int i,int j){
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

 

 

퀵 정렬을 조금 더 공부해야겠다..