코딩테스트 불길한 수열 문제

문제

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 1

8
7
5
4
3
1

n번째 최소값/최대값

최소 힙이나 최대 힙을 사용하면 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]

Share: