코딩테스트 불길한 수열 문제
문제
4, 13, 413, 134와 같이 숫자가 13과 4를 이용해서 만들 수 있는 수를 불길한 수라고 한다. 그리고 불길한 수가 오름차순으로 나열된 수열을 불길한 수열이라고 한다. 예를 들어 불길한 수열은 다음과 같이 나열될 수 있다.
ex) S = {4, 13, 44, 134, 413, 444, 1313......}
n이 매개변수로 주어질 때, 불길한 수열에서 n번째 불길한 수를 구하시오.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | from heapq import heappush, heappop, heapify def solution2(n): q = [4,13] heapify(q) while n: n-=1 v=heappop(q) heappush(q,v*10+4) heappush(q,v*100+13) return v print(solution2(100)) | cs |
핵심은 heap을 사용한다.
heapq 모듈은 이진트리 기반의 최소 힙 자료구조를 제공한다.
min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, min heap에서
가장 작은 원소는 언제나 인덱스 0이다.
내부적으로 min heap 내의 모든 원소(k)는 2k+1, 2k+2보다 작거나 같다.
최소 힙 생성
heapq 모듈은 파이썬의 보통 리스트를 마치 최소 힙처럼 다룰 수 있도록 도와준다.
그래서 빈 리스트를 생성한 다음에 heapq 모듈의 함수를 호출할 때마다 이 리스트를 인자로 넘긴다.
heap = []from heapq import heappush
heappush(heap, 4)
heappush(heap, 1)
heappush(heap, 7)
heappush(heap, 3)
print(heap)[1, 3, 7, 4]결과값처럼 가장 작은 1이 인덱스 0에 위치하며, 인덱스 1(=k)에 위치한 3은 인덱스
3(=2k+1)에 위치한 4보다 크므로 힙의 공식을 만족합니다.
내부적으로 이진트리에 원소를 추가하는 heappush() 함수는 O(logN)의 시간복잡도를 가진다.
힙에서 원소 삭제
from heapq import heappop
print(heappop(heap))
print(heap)1
[3, 4, 7]heapq 모듈의 heappop() 함수를 이용하여 힙에서 원소를 삭제할 수 있다.
원소를 삭제할 대상 리스트를 인자로 넘기면, 가장 작은 원소를 삭제 후에 그 값을 리턴한다.
힙에서의 원소 삭제 역시 O(logN)의 시간복잡도를 가진다.
기존 리스트를 힙으로 변환
from heapq import heapify
heap = [4, 1, 7, 3, 8, 5]
heapify(heap)
print(heap)[1, 3, 5, 4, 8, 7]heapq 모듈의 heapify()라는 함수를 사용하면 리스트를 힙으로 만들 수 있다.
heapify() 함수의 성능은 O(N)이다.
nums = [4, 1, 7, 3, 8, 5]
heap = nums[:]
heapify(heap)
print(nums)
print(heap)[4, 1, 7, 3, 8, 5]
[1, 3, 5, 4, 8, 7]이때, 주의할 점은 새로운 리스트를 반환하는 것이 아니라 인자로 넘긴 리스트를 직접 변경한다.
최대 힙
heapq 모듈은 최소 힙을 기본으로 동작하기 때문에 최대 힙으로 활용하려면 약간의 요령이 필요하다.
바로 힙에 튜플을 원소로 추가하거나 삭제하면, 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 원리를 이용하면된다.
다시 말해, (우선순위, 값) 구조의 튜플을 힙에 추가하거나 삭제하면 된다.
그리고 힙에서 값을 읽어올 때는 각 튜플에서 인덱승 1에 있는 값을 취하면된다.
from heapq import heappush, heappop
nums = [4, 1, 7, 3, 8, 5]
heap = []
for num in nums:
heappush(heap, (-num, num)) # (우선 순위, 값)
while heap:
print(heappop(heap)[1]) # index 18
7
5
4
3
1n번째 최소값/최대값
최소 힙이나 최대 힙을 사용하면 n번째로 작은 값이나 큰 값을 효과적으로 구할 수 있다.
from heapq import heappush, heappop
def nth_smallest(nums, n):
heap = []
for num in nums:
heappush(heap, num)
nth_min = None
for _ in range(n):
nth_min = heappop(heap)
return nth_min
print(nth_smallest([4, 1, 7, 3, 8, 5], 3))4혹은
from heapq import nsmallest
nsmallest(3, [4, 1, 7, 3, 8, 5])[-1]nsmallest() 함수를 사용한다.
nsmallest() 함수는 주어진 리스트에서 가장 작은 n개의 값을 담은 리스트를 반환한다.
그 리스트의 마지막 값이 n번째 작은 값이 된다.
반대로 n번째 최대값을 구할 때는 nlargest()함수를 사용한다.
from heapq import nlargest
nlargest(3, [4, 1, 7, 3, 8, 5])[-1]
