정글사관학교 16일차 TIL: 힙 자료구조, heapq 모듈, 가운데를 말해요

0. 목차

1. 힙 자료구조 & heapq: Heap 모듈 in python

2. 가운데를 말해요

 

 

원래 오늘 정리하고 싶었던 문제는 원 영역이었는데..아 진짜 다른 풀이를 봐도 뭔가 명쾌한 이거다! 싶은 느낌이 오지 않는다. 오일러 지표(v-e+f=c)를 푸는 풀이, DFS를 이용하는 풀이 등 다양한 방식이 있는데 아직 잘 체감이 되지 않는다. 이 문제는 손에 좀 더 익히고 나서 정리하도록 할 예정.

 

오늘은 힙 구조에 대해 정리하고 관련 문제인 "가운데를 말해요"에 대해 정리한다. 내일은 이제까지 미처 정리 못했던 문제 중 중요한 애들(이분 탐색 - 사냥꾼 / 스택 - 원 영역, 괄호의 값, 크게 만들기 / 큐 - 뱀) 다시 정리해볼 예정..!

 

할 수 있다!

---

1. 힙 자료구조 & Heap 모듈

힙 자료구조에 대한 정리는 이전 글에서 간략하게나마 정리해둔 게 있다. 오늘은 여기에 추가하여 개념 정리를 적고 나아가 파이썬에서 힙 구조를 구현하는 방법으로써 heapq 모듈을 소개할 것이다.

 

이진 트리(Binary Tree)

• 데이터를 표현할 때 각 노드에 데이터를 담은 뒤에 노드를 두 개씩 이어 붙이는 구조. 트리 구조에 맞게 부모 노드에서 자식 노드로 가지가 뻗힌다.

힙(Heap)

• 최솟값/최댓값을 빠르게 찾아내기 위해 완전 이진 트리를 기반으로 하는 트리. O(nlogn)의 시간복잡도로 최대 힙 정렬 가능.

Heapq 모듈

Heapq는 파이썬 내장 모듈로, 이진 트리 기반의 최소 힙 자료구조를 제공한다. 최소 힙 자료구조로 고정되어 있기 때문에 가장 작은 값은 이진 트리 루트에 위치한다.

 

내부적으로 최소 힙 내 모든 원소 k는 항상 자식 원소를 (2k+1, 2k+2)보다 크기가 작거나 같도록 원소를 추가/삭제한다.

 

모듈을 import하는 방법

import heapq

 

최소 힙을 생성하는 방식

heapq는 파이썬 리스트를 힙처럼 쓸 수 있게 하는 방식이다. 따라서 따로 자료구조 형태를 제공하지는 않고 빈 리스트에 heap 모듈을 호출할 때마다 리스트를 인자로 넘기는 식이다.

 

힙에 원소를 추가하는 방법

heap 모듈 내 heappush() 함수를 이용한다.

 

import heapq

heap = list() # 빈 리스트 하나 불러옴

heapq.heappush(heap, 4) # (원소 추가할 대상 리스트, 해당 리스트에 추가할 원소)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3)

# heapq = [1, 3, 7, 4]

heappush() 메소드는 O(logN)의 시간복잡도를 지닌다. (최소힙을 찾기까지 절반씩 나눠가면서 찾기 때문 => 여기 링크에서 이분 탐색의 시간 복잡도를 검색하면 이해 가능)

 

힙에서 원소를 삭제하는 방법

heap.heappop()을 이용하면 가장 작은 원소를 삭제 후 이 값을 리턴한다. 이후에는 그 다음으로 최소인 값을 힙 리스트의 0번 인덱스에 위치시킨다. 이 역시 시간복잡도는 O(logn).

 

최솟값 삭제하지 않고 얻기

힙에서 최솟값 빼내지 말고 읽기만 하고 싶다면? heap[0]를 출력한다.

 

*주의!

인덱스 0에 가장 작은 원소가 있다고 해서 인덱스 1에 두번째로 작은 원소, 인덱스 2에 세번째로 작은 원소가 있다는 보장이 없다. heappop()을 호출해서 원소를 삭제할 때마다 이진 트리를 재배치해서 새로운 최솟값을 인덱스 0에 위치시키기 때문.

따라서 두번째로 작은 원소를 얻으려면? heap[1]을 출력하는 게 아니라 heappop()을 이용해 가장 첫번째로 작은 원소를 삭제한 뒤 heap[0]으로 새로운 최솟값에 접근한다.

기존 리스트를 힙으로 변환

이미 원소가 들어있는 리스트를 힙으로 변환하려면 heap.heapify() 모듈을 사용한다.

 

import heapq

heap = [4, 1, 7, 3, 8, 5]
heapq.heapify(heap) # 바꾸려는 리스트
# heap = [1, 3, 5, 4, 8, 7]

리스트 내부의 원소 위에서 다룬 힙구조에 맞게 재배치한다. 이는 비어있는 리스트를 생성한 뒤 heappush()함수로 원소를 하나씩 추가한 것과 동일한 방식이다. 원소를 하나씩 배치하기 때문에 O(n)의 시간복잡도를 지닌다.

 

heapq로 최대 힙 만들기

위에서 설명했듯, heapq는 기본적으로 최소 힙만을 제공한다. 하지만 약간의 트릭을 활용하면 최대 힙을 구현할 수 있다. 바로 힙에 튜플을 원소로 추가하는 방식이다.

 

튜플에다가 (우선순위, 값)을 각각 입력하는데 이때

 

우선순위 = -값

형태로 만들면 우선순위를 기준으로 했을 때 순서가 반대로 된다. 실제로 읽어올 때는 튜플의 인덱스 1(실제값)을 읽어오고, 힙 구조로 만들 때는 인덱스 0(우선순위)를 이용해 -값을 기준으로 최소 힙이 구성되는 원리를 활용한다.

 

따라서 최대 힙을 구현하려면 각 값에 대한 우선순위를 구한 다음 (우선순위, 값) 구조의 튜플을 힙에 추가하거나 삭제한다. 힙에서 값을 읽어올 때는 각 튜플에서 인덱스에 있는 값만 취하고 우선순위는 생각하지 않는다.

 

코드를 보자.

import heapq
nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
    heapq.heappush(heap, (-num, num)) # -> 원래는 넣고자 하는 값 num만 넣었는데 이를 튜플로 넣는다.

while heap:
    print(heapq.heappop(heap[1]) # 튜플에서 인덱스 1에 해당하는 값을 print = 최대를 가져온다.

참고 자료 1

참고 자료 2

 

 

 

 

2. 가운데를 말해요(우선순위 큐)

 

문제 링크는 여기(#1655)

 

1번 풀이: sort 이용 - 시간 복잡도: O(n^2)

 

안 풀린다는 걸 알면서도 자연스럽게 손이 가요 손이 가는 sort...어차피 힙 써야 하는 건 알고 있다만 뒤에 디테일한 구현할 때 아이디어 참조할 겸 빠르게 짜봤다. 당연히 결과는 시간 초과.

 

# 1. sort 이용 => 시간 초과.. 알면서도 그냥 해봄..

from sys import stdin

n = int(stdin.readline())
ans = []
for i in range(n):
    num = int(stdin.readline())
    ans.append(num)
    ans.sort()
    if i % 2 == 0:
        print(ans[i//2])
    else:
        if ans[i // 2] < ans[i //2 +1]:
            print(ans[i//2])
        else:
            print(ans[i // 2+1])

 

 

여기서 착안한 아이디어는

 

1. 숫자를 꺼낼 때마다 그 값을 리스트에 저장한 뒤, 해당 숫자의 인덱스의 절반 앞까지는 다 버린다.

2. 절반을 기점으로, 리스트 내 원소 수가 홀수면 바로 절반에 해당하는 인덱스를 정답으로 출력하면 되고(인덱스는 0부터 세기 때문에 원소 수가 홀수면 마지막 인덱스 i는 짝수에 해당)  원소 개수가 짝수면 중간에 있는 두 수를 비교한 다음 작은 수를 출력하는 쪽으로 코드를 짰다. 정답은 나오는 것 같은데 당연히 시간초과.

 

2. 힙을 이용한 풀이: 이중 for문 써서 힙 쓴 의미가 없음...시간 초과

 

이 개념을 힙으로 연결해서 다시 짜봤다. 사실 힙을 제대로 쓴 풀이가 아니긴 한데, heap을 갖다 쓰긴 했지만 heap으로 push하기까지 for문을 두 번 돌기 때문에 시간복잡도는 여전히 O(n^2)에 해당한다. 아래 코드를 보자.

 

# 힙 정렬하면? => 얘도 시간초과.. for문 두 번 돌아서 의미가 없는듯
from sys import stdin
import heapq

n = int(stdin.readline())

ans = []
for i in range(n):
    num = int(stdin.readline())
    heapq.heappush(ans, num)
    temp = ans[:]
    for _ in range(i // 2):
        heapq.heappop(temp)
    if i % 2 == 0:
        pass
    else:
        a = heapq.heappop(temp)
        b = heapq.heappop(temp)
        if a < b:
            print(a)
        else:
            print(b)
        continue
    mid = heapq.heappop(temp)
    print(mid)

 

3번 풀이: heap을 이용한 풀이 2 (풀이 참조 링크)

 

위의 풀이와 달리 시간을 줄일 수 있었던 이유는 아래 코드에서는 for문을 한 번만 쓰고 대신에 두 개의 heap을 이용해 O(NlogN)을 유지했다는 점 때문이다.

 

자세한 풀이는 링크를 참조.

 

시간이 벌써 세시라 오늘은 여기까지..

from sys import stdin
import heapq

n = int(stdin.readline())

leftheap = []
rightheap = []
answer = []

for i in range(n):
    input_num = int(stdin.readline())

    if len(leftheap) == len(rightheap):
        heapq.heappush(leftheap, (-input_num, input_num)) # 양쪽에 같은 개수가 들어가면 (==
    else:
        heapq.heappush(rightheap, (input_num, input_num))

    if rightheap and leftheap[0][1] > rightheap[0][0]: # rightheap이 있고 & leftheap[0][1]
        min = heapq.heappop(rightheap)[0]
        max = heapq.heappop(leftheap)[1]
        heapq.heappush(leftheap, (-min, min))
        heapq.heappush(rightheap, (max, max))

    answer.append(leftheap[0][1])
for j in answer:
    print(j)