문제 링크

문제

N개의 서로 다른 양의 정수가 저장된 배열 A가 있다. 선택 정렬로 배열 A를 오름차순 정렬할 경우 K 번째 교환되는 수를 구하자.

코드

처음에 문제를 제대로 읽지 않고, 이전에 풀었던 28116번 방법을 적용하려다가 낭패를 겪었다.

28116번 문제는 수열 A는 무조건 1부터 N까지의 숫자 중 하나씩 등장하지만, 이 문제에서는 아니다. 양의 정수 A의 범위는 $1 <= A_i <= 10^9$다. 때문에 배열만으로는 문제를 풀 수 없어서 TreeMap이라는 자료구조를 사용했다.

TreeMap은 이진트리를 기반으로 한 Map 컬렉션으로, 새로운 데이터를 추가할 때마다 정렬이 된다. 그래서 이를 사용해 쉽게 구현할 수 있었다. TreeMap 생성 시 내림차순으로 정렬하도록 설정해줬다.

C언어로 구현한 사람들의 코드를 보니 정렬 부분을 직접 구현한 것 같던데 추후 성능 좋을 정렬 알고리즘을 알게 된다면 여기에 적용해보고 싶다.

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

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        StringTokenizer st;

        st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());

        // 각 숫자가 몇 번째 인덱스에 있는지 저장하고 있는 TreeMap 객체
        TreeMap<Integer, Integer> idxMap = new TreeMap<>(Collections.reverseOrder());

        // 수열 A를 저장하고 있는 배열
        int[] numArr = new int[n + 1];

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            int num = Integer.parseInt(st.nextToken());
            idxMap.put(num, i);
            numArr[i] = num;
        }

        int times = 0;      // 교환 횟수
        int i = n;
        
        // TreeMap 객체 생성 시 내림차순으로 정렬되게 설정하였으므로 maxValue는 큰 수부터 오게 된다.
        for (Integer maxValue : idxMap.keySet()) {
            // 가장 큰 값이 저장된 인덱스
            int maxIndex = idxMap.get(maxValue);

            // i번째 데이터가 최대값이 아니라면...
            if (numArr[i] != maxValue) {
                // swap이 일어날 때만 카운트
                if (++times >= k) {
                    // 항상 큰 값이 뒤로 오도록 정답 출력
                    sb.append(numArr[i]).append(" ").append(maxValue);
                    System.out.println(sb);
                    return;
                }

                // TreeMap과 배열 값 각각 swap
                idxMap.put(numArr[i], maxIndex);
                numArr[maxIndex] = numArr[i];
            }

            i--;
        }

        System.out.println(-1);
    }
}