[python/알고리즘] 계수 정렬 | 백준 10989 수 정렬하기3

백준 10989 - 수 정렬하기3

문제

N개의 수가 주어졌을 때, 이를 오름차순으로 정렬하는 프로그램을 작성하시오.
첫째 줄에 수의 개수 N(1 ≤ N ≤ 10,000,000)이 주어진다. 둘째 줄부터 N개의 줄에는 수가 주어진다. 이 수는 10,000보다 작거나 같은 자연수이다.
시간 제한: 5초
메모리 제한: 8MB

접근

이 문제는 단순히 정렬하는 것이 중요한 것이 아니다. 일반 정렬을 사용하면 시간 복잡도가 크기 때문에 메모리 제한을 피하기 위해서는 시간 복잡도가 O(n)인 계수 정렬(Counting sort) 알고리즘을 이용해야 한다.

코드

import sys

N = int(input())

count = [0] * 10001

for _ in range(N):
    num = int(sys.stdin.readline())
    count[num] += 1

for i in range(1, 10001):
    if count[i] != 0:
        for _ in range(count[i]):
            print(i)

해설

계수 정렬은 각 숫자의 개수를 세는 방식으로 진행된다. 숫자의 범위가 제한적이고 (ex. 10000보다 작거나 같은 수) 큰 수의 데이터를 정렬해야 할 때 매우 효율적이다.

우선 0으로 채워진 빈 배열 count를 생성한다. 이때 count의 길이는 문제에 등장할 수 있는 가장 큰 수 + 1의 크기여야 한다.

첫 번째 for문에서는 입력을 받는다. 입력 받은 수를 인덱스로 하는 자리에 1을 더한다. 만약 등장하는 가장 큰 수가 5이고, 입력 값이 3, 4, 5, 3이라면 첫 번째 for 루프를 순회한 후 count는 [0, 0, 0, 2, 1, 1]일 것이다.

두 번째 for문에서는 1부터 max값인 10000까지 count의 인덱스를 돌면서 값이 0이 아닌 인덱스를 만나면 그 인덱스 'i'를 그 자리에 저장된 값만큼 출력한다. 해당 숫자 'i'가 그 숫자만큼 등장했다는 의미이기 때문이다. 앞선 예의 count [0, 0, 0, 2, 1, 1]을 순회하다가 0이 아닌 값 2가 등장하면 그 인덱스인 3을 두 번 출력하고, 4를 한 번, 4를 한 번 출력해서 정렬된 [3, 3, 4, 5]가 출력된다.