[Java] 11004 K번째 수

문제

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

 

입력

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

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

 

출력

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

 

import java.io.*;
import java.util.*;

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 num = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());
        st = new StringTokenizer(br.readLine());
        int[] arr = new int[num];
        
        for(int i=0; i<num; i++){
            arr[i] = Integer.parseInt(st.nextToken());
        }
        quickSort(arr, 0, num-1, k-1);
        System.out.println(arr[k-1]);
    }
    
    public static void quickSort(int[] arr, int start, int end, int k){
        if(start < end){
            int pivot = partition(arr, start, end);
            if(pivot == k)
                return;
            else if(k < pivot)
                quickSort(arr, start, pivot-1, k);
            else
                quickSort(arr, pivot+1, end, k);
        }
    }
    
    public static int partition(int[] arr, int start, int end){
        if(start+1 == end){
            if(arr[start] > arr[end])
                swap(arr, start, end);
            return end;
        }
        int middle = (start + end) / 2;
        swap(arr, start, middle);
        int pivot = arr[start];
        int i = start+1, j = end;
        
        while(i <= j){
            while(pivot < arr[j] && j > 0){
                j--;
            }
            while(pivot > arr[i] && i < arr.length-1){
                i++;
            }
            if(i <= j){
                swap(arr, i++, j--);
            }
        }
        arr[start] = arr[j];
        arr[j] = pivot;
        return j;
    }
    
    public static void swap(int[] arr, int i, int j){
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }
}

▶ int num = Integer.parseInt(st.nextToken()) : 데이터의 개수를 입력받아 변수 num에 저장함

▶ int k = Integer.parseInt(st.nextToken()) : k번째 수를 입력받아 변수 k에 저장함

▶ int[] arr = new int[num] : num 크기만큼의 배열 arr 선언

▶ for(int i=0; i<num; i++){ arr[i] = Integer.parseInt(st.nextToken()) } : 배열의 값을 입력 받음

 

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

▶ start 값이 end 값보다 작고, k번째 수가 pivot이면 더이상 구할 필요가 없어 return 함

▶ k가 pivot보다 작으면 왼쪽 그룹만 정렬함

▶ k가 pivot보다 크면 오른쪽 그룹만 정렬함

 

public static int partition(int[] arr, int start, int end){
    if(start+1 == end){
        if(arr[start] > arr[end])
            swap(arr, start, end);
        return end;
    }
    int middle = (start + end) / 2;
    swap(arr, start, middle);
    int pivot = arr[start];
    int i = start+1, j = end;

    while(i <= j){
        while(pivot < arr[j] && j > 0){
            j--;
        }
        while(pivot > arr[i] && i < arr.length-1){
            i++;
        }
        if(i <= j){
            swap(arr, i++, j--);
        }
    }
    arr[start] = arr[j];
    arr[j] = pivot;
    return j;
}

▶ start 인덱스의 값이 end 인덱스의 값보다 큰 경우 end 값을 return 함

▶ ( start 값 + end 값 ) / 2로 중앙값을 계산함

▶ 중앙값을 start와 위치를 변경함

▶ pivot을 시작 위치 값 arr[start]로 저장함

▶ i는 시작점, j는 종료점으로 지정함

▶ 피벗보다 작은 수가 나올 때까지 j--

▶ 피벗보다 큰 수가 나올 때까지 i++

▶ i == j pivot의 값을 양쪽으로 분리한 가운데에 오도록 설정함