[Python] 백준 23883 알고리즘 수업 - 선택 정렬 3

문제

오늘도 서준이는 선택 정렬 수업 조교를 하고 있다. 아빠가 수업한 내용을 학생들이 잘 이해했는지 문제를 통해서 확인해보자.

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

N이 매우 커서 시간 초과를 고민하고 있는 우리 서준이를 도와주자.

크기가 N인 배열에 대한 선택 정렬 의사 코드는 다음과 같다.

selection_sort(A[1..N]) { # A[1..N]을 오름차순 정렬한다.
    for last <- N downto 2 {
        A[1..last]중 가장 큰 수 A[i]를 찾는다
        if (last != i) then A[last] <-> A[i]  # last와 i가 서로 다르면 A[last]와 A[i]를 교환
    }
}

입력

첫째 줄에 배열 A의 크기 N(5 ≤ N ≤ 500,000), 교환 횟수 K(1 ≤ K ≤ N)가 주어진다.

다음 줄에 서로 다른 배열 A의 원소 A1, A2, ..., AN이 주어진다. (1 ≤ Ai ≤ 109)

출력

K 번째 교환되는 두 개의 수를 작은 수부터 한 줄에 출력한다. 교환 횟수가 K 보다 작으면 -1을 출력한다.

예제 입력 1 복사

5 2
3 1 2 5 4

예제 출력 1 복사

2 3

3 1 2 5 4 (4와 5가 교환됨) -> 3 1 2 4 5 (A[1..4]에서 4가 가장 크므로 교환이 발생하지 않음) -> 3 1 2 4 5 (2와 3이 교환됨) -> 2 1 3 4 5 (1과 2가 교환됨) -> 1 2 3 4 5. 총 3회 교환이 발생하고 두 번째 교환은 2와 3이다.

예제 입력 2 복사

5 4
3 1 2 5 4

예제 출력 2 복사

-1

교환 횟수 3이 K 보다 작으므로 -1을 출력한다.

 

---

코드

import sys

def selection_sort(N, K, A):
    # 정렬된 리스트
    sortedA = sorted(A)
    numDict = {}
    
    # 딕셔너리  key : A원소 value : 정렬 전 인덱스
    for i, j in enumerate(A):
      numDict[j] = i
      
    cnt = 0
    for i in range(N-1, -1, -1): # i : 현재 정렬할 인덱스
        if sortedA[i] != A[i]: # 현재 인덱스가 아니면?
          result = [A[i], sortedA[i]]
          A[i], A[numDict[sortedA[i]]] = A[numDict[sortedA[i]]], A[i]
          numDict[result[0]], numDict[result[1]] = numDict[result[1]], numDict[result[0]]
          cnt += 1
        if K == cnt : print(*result);return
    if K > cnt : print(-1);return


N, K = map(int, sys.stdin.readline().rstrip().split())
A = list(map(int, sys.stdin.readline().rstrip().split()))
selection_sort(N, K, A)

---

리뷰

단순히 순서대로 선택정렬을 하였을 때 배열의 크기로 인해 시간 초과가 되는 문제였다.

 

가장 먼저 시간 복잡도 부분을 해결하기 위해서 어떤 것을 사용해볼까? 라는 질문에

현재 정렬 세션의 가장 큰 값을 찾는 부분을 힙 정렬을 사용하는 것을 시도하였다.

이를 위해서 heapq 모듈을 사용하려 했지만

매 반복 대상 부분리스트마다 힙 정렬을 위해 heapify를 해야하는데

이를 위해서 계속 따로 힙을 생성하고 최대값을 찾고, 해당 값의 인덱스를 다시 찾아야하는 점이

비효율적일 것 같아서 다른 방법을 찾아보았다.

 

그래서 두 번째로 고려해본 것은 딕셔너리를 이용하는 것이었다.

딕셔너리 자료형은 내부적으로 해시 테이블을 이용하므로 데이터의 '조회 및 수정이 O(1)'의 시간복잡도를 가지기 때문에

이를 이용하면 적합하지 않을까? 라고 생각하였다.

 

1) 정렬 리스트 생성과 딕셔너리 생성

 # 정렬된 리스트
sortedA = sorted(A)
numDict = {}

# 딕셔너리  key : A원소 value : 정렬전 인덱스
for i, j in enumerate(A):
  numDict[j] = i

먼저 최악의 경우에도 nlogn의 시간을 보장하는 Timsort알고리즘으로 구현된 sorted()를 이용하여

주어진 리스트 A를 정렬한리스트 sortedA를 생성하였다.

 

enumerate()를 이용하여 numDIct 딕셔너리에

key = A의 원소, value=해당 원소의 인덱스

로 삽입하였다.

 

2) 선택정렬 & 결과값 찾기

for i in range(N-1, -1, -1): # i : 현재 정렬할 인덱스
    if sortedA[i] != A[i]: # 현재 인덱스가 아니면?
      result = [A[i], sortedA[i]]
      A[i], A[numDict[sortedA[i]]] = A[numDict[sortedA[i]]], A[i]
      numDict[result[0]], numDict[result[1]] = numDict[result[1]], numDict[result[0]]
      cnt += 1
    if K == cnt : print(*result);return
if K > cnt : print(-1);return

sortedA와 A의 마지막 원소부터 비교하며 정렬을 진행한다.

만약 현재 인덱스의 값이 서로 다르다면 A는 정렬되어야한다.

 

정렬할 값의 인덱스는 "numDict[sortedA[현재인덱스]]"에 담겨 있다.

이를 이용하여 A[현재 인덱스] 와 A[numDict[sortedA[현재인덱스]]] 를 서로 바꾸어 정렬한다.

 

바꾼 것에 대한 정보를 numDict에도 적용한다.