10_heap.py - Examples

  1"""
  2힙과 우선순위 큐 (Heap & Priority Queue)
  3Heap and Priority Queue
  4
  5힙 자료구조와 관련 알고리즘을 구현합니다.
  6"""
  7
  8from typing import List, Optional, Tuple, Any
  9import heapq
 10
 11
 12# =============================================================================
 13# 1. 최소 힙 (Min Heap) 직접 구현
 14# =============================================================================
 15
 16class MinHeap:
 17    """
 18    최소 힙 (완전 이진 트리)
 19    - 부모 노드 ≤ 자식 노드
 20    - 배열로 구현: parent(i) = (i-1)//2, left(i) = 2i+1, right(i) = 2i+2
 21    """
 22
 23    def __init__(self):
 24        self.heap: List[int] = []
 25
 26    def __len__(self) -> int:
 27        return len(self.heap)
 28
 29    def __bool__(self) -> bool:
 30        return len(self.heap) > 0
 31
 32    def _parent(self, i: int) -> int:
 33        return (i - 1) // 2
 34
 35    def _left(self, i: int) -> int:
 36        return 2 * i + 1
 37
 38    def _right(self, i: int) -> int:
 39        return 2 * i + 2
 40
 41    def _swap(self, i: int, j: int) -> None:
 42        self.heap[i], self.heap[j] = self.heap[j], self.heap[i]
 43
 44    def _sift_up(self, i: int) -> None:
 45        """삽입 후 위로 재정렬 - O(log n)"""
 46        while i > 0:
 47            parent = self._parent(i)
 48            if self.heap[i] < self.heap[parent]:
 49                self._swap(i, parent)
 50                i = parent
 51            else:
 52                break
 53
 54    def _sift_down(self, i: int) -> None:
 55        """삭제 후 아래로 재정렬 - O(log n)"""
 56        n = len(self.heap)
 57
 58        while True:
 59            smallest = i
 60            left = self._left(i)
 61            right = self._right(i)
 62
 63            if left < n and self.heap[left] < self.heap[smallest]:
 64                smallest = left
 65            if right < n and self.heap[right] < self.heap[smallest]:
 66                smallest = right
 67
 68            if smallest != i:
 69                self._swap(i, smallest)
 70                i = smallest
 71            else:
 72                break
 73
 74    def push(self, val: int) -> None:
 75        """삽입 - O(log n)"""
 76        self.heap.append(val)
 77        self._sift_up(len(self.heap) - 1)
 78
 79    def pop(self) -> int:
 80        """최솟값 제거 및 반환 - O(log n)"""
 81        if not self.heap:
 82            raise IndexError("pop from empty heap")
 83
 84        result = self.heap[0]
 85        last = self.heap.pop()
 86
 87        if self.heap:
 88            self.heap[0] = last
 89            self._sift_down(0)
 90
 91        return result
 92
 93    def peek(self) -> int:
 94        """최솟값 조회 - O(1)"""
 95        if not self.heap:
 96            raise IndexError("peek from empty heap")
 97        return self.heap[0]
 98
 99    def heapify(self, arr: List[int]) -> None:
100        """배열을 힙으로 변환 - O(n)"""
101        self.heap = arr[:]
102        # 마지막 비-리프 노드부터 sift_down
103        for i in range(len(self.heap) // 2 - 1, -1, -1):
104            self._sift_down(i)
105
106
107# =============================================================================
108# 2. 최대 힙 (Max Heap)
109# =============================================================================
110
111class MaxHeap:
112    """최대 힙 - 부모 노드 ≥ 자식 노드"""
113
114    def __init__(self):
115        self.heap: List[int] = []
116
117    def __len__(self) -> int:
118        return len(self.heap)
119
120    def push(self, val: int) -> None:
121        """삽입 - O(log n)"""
122        # 최소 힙에 음수로 저장
123        heapq.heappush(self.heap, -val)
124
125    def pop(self) -> int:
126        """최댓값 제거 및 반환 - O(log n)"""
127        return -heapq.heappop(self.heap)
128
129    def peek(self) -> int:
130        """최댓값 조회 - O(1)"""
131        return -self.heap[0]
132
133
134# =============================================================================
135# 3. 힙 정렬 (Heap Sort)
136# =============================================================================
137
138def heap_sort(arr: List[int]) -> List[int]:
139    """
140    힙 정렬
141    시간복잡도: O(n log n)
142    공간복잡도: O(1) (in-place)
143    """
144    n = len(arr)
145    result = arr[:]
146
147    def sift_down(arr: List[int], n: int, i: int) -> None:
148        while True:
149            largest = i
150            left = 2 * i + 1
151            right = 2 * i + 2
152
153            if left < n and arr[left] > arr[largest]:
154                largest = left
155            if right < n and arr[right] > arr[largest]:
156                largest = right
157
158            if largest != i:
159                arr[i], arr[largest] = arr[largest], arr[i]
160                i = largest
161            else:
162                break
163
164    # 1. 최대 힙 구성 - O(n)
165    for i in range(n // 2 - 1, -1, -1):
166        sift_down(result, n, i)
167
168    # 2. 하나씩 추출 - O(n log n)
169    for i in range(n - 1, 0, -1):
170        result[0], result[i] = result[i], result[0]
171        sift_down(result, i, 0)
172
173    return result
174
175
176# =============================================================================
177# 4. K번째 요소 찾기
178# =============================================================================
179
180def kth_largest(arr: List[int], k: int) -> int:
181    """
182    K번째로 큰 요소 - 최소 힙 사용
183    시간복잡도: O(n log k)
184    공간복잡도: O(k)
185    """
186    min_heap = []
187
188    for num in arr:
189        heapq.heappush(min_heap, num)
190        if len(min_heap) > k:
191            heapq.heappop(min_heap)
192
193    return min_heap[0]
194
195
196def kth_smallest(arr: List[int], k: int) -> int:
197    """
198    K번째로 작은 요소 - 최대 힙 사용
199    시간복잡도: O(n log k)
200    공간복잡도: O(k)
201    """
202    max_heap = []
203
204    for num in arr:
205        heapq.heappush(max_heap, -num)
206        if len(max_heap) > k:
207            heapq.heappop(max_heap)
208
209    return -max_heap[0]
210
211
212# =============================================================================
213# 5. 중앙값 스트림 (Median Finder)
214# =============================================================================
215
216class MedianFinder:
217    """
218    데이터 스트림의 중앙값
219    - 최대 힙 (왼쪽 절반): 작은 값들
220    - 최소 힙 (오른쪽 절반): 큰 값들
221    """
222
223    def __init__(self):
224        self.max_heap: List[int] = []  # 왼쪽 (작은 값들)
225        self.min_heap: List[int] = []  # 오른쪽 (큰 값들)
226
227    def add_num(self, num: int) -> None:
228        """숫자 추가 - O(log n)"""
229        # 왼쪽 힙에 추가
230        heapq.heappush(self.max_heap, -num)
231
232        # 왼쪽 최댓값을 오른쪽으로 이동
233        heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
234
235        # 크기 균형 유지 (왼쪽 ≥ 오른쪽)
236        if len(self.min_heap) > len(self.max_heap):
237            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))
238
239    def find_median(self) -> float:
240        """중앙값 반환 - O(1)"""
241        if len(self.max_heap) > len(self.min_heap):
242            return -self.max_heap[0]
243        return (-self.max_heap[0] + self.min_heap[0]) / 2
244
245
246# =============================================================================
247# 6. 우선순위 큐 응용: 작업 스케줄링
248# =============================================================================
249
250def schedule_tasks(tasks: List[Tuple[int, int, str]]) -> List[str]:
251    """
252    우선순위 기반 작업 스케줄링
253    tasks: [(우선순위, 도착시간, 작업명), ...]
254    낮은 우선순위 번호 = 높은 우선순위
255    """
256    # (우선순위, 도착시간, 작업명) 형태로 힙에 추가
257    task_heap = []
258    for priority, arrival, name in tasks:
259        heapq.heappush(task_heap, (priority, arrival, name))
260
261    schedule = []
262    while task_heap:
263        _, _, name = heapq.heappop(task_heap)
264        schedule.append(name)
265
266    return schedule
267
268
269# =============================================================================
270# 7. K개 정렬된 리스트 병합
271# =============================================================================
272
273def merge_k_sorted_lists(lists: List[List[int]]) -> List[int]:
274    """
275    K개의 정렬된 리스트 병합
276    시간복잡도: O(N log k), N = 전체 원소 수, k = 리스트 수
277    """
278    result = []
279    min_heap = []
280
281    # 각 리스트의 첫 번째 원소를 힙에 추가
282    for i, lst in enumerate(lists):
283        if lst:
284            heapq.heappush(min_heap, (lst[0], i, 0))  # (값, 리스트 인덱스, 원소 인덱스)
285
286    while min_heap:
287        val, list_idx, elem_idx = heapq.heappop(min_heap)
288        result.append(val)
289
290        # 다음 원소가 있으면 힙에 추가
291        if elem_idx + 1 < len(lists[list_idx]):
292            next_val = lists[list_idx][elem_idx + 1]
293            heapq.heappush(min_heap, (next_val, list_idx, elem_idx + 1))
294
295    return result
296
297
298# =============================================================================
299# 8. Top K 빈도 요소
300# =============================================================================
301
302def top_k_frequent(nums: List[int], k: int) -> List[int]:
303    """
304    가장 빈번한 K개 요소
305    시간복잡도: O(n log k)
306    """
307    from collections import Counter
308
309    freq = Counter(nums)
310
311    # 최소 힙으로 K개만 유지
312    min_heap = []
313    for num, count in freq.items():
314        heapq.heappush(min_heap, (count, num))
315        if len(min_heap) > k:
316            heapq.heappop(min_heap)
317
318    return [num for count, num in min_heap]
319
320
321# =============================================================================
322# 9. 가장 가까운 K개 점
323# =============================================================================
324
325def k_closest_points(points: List[Tuple[int, int]], k: int) -> List[Tuple[int, int]]:
326    """
327    원점에서 가장 가까운 K개 점
328    시간복잡도: O(n log k)
329    """
330    # 최대 힙으로 K개만 유지 (거리의 음수 저장)
331    max_heap = []
332
333    for x, y in points:
334        dist = x * x + y * y
335        heapq.heappush(max_heap, (-dist, x, y))
336        if len(max_heap) > k:
337            heapq.heappop(max_heap)
338
339    return [(x, y) for _, x, y in max_heap]
340
341
342# =============================================================================
343# 10. 회의실 문제 (Meeting Rooms)
344# =============================================================================
345
346def min_meeting_rooms(intervals: List[Tuple[int, int]]) -> int:
347    """
348    필요한 최소 회의실 수
349    intervals: [(시작시간, 종료시간), ...]
350    시간복잡도: O(n log n)
351    """
352    if not intervals:
353        return 0
354
355    # 시작 시간으로 정렬
356    intervals.sort(key=lambda x: x[0])
357
358    # 종료 시간을 저장하는 최소 힙
359    end_times = []
360    heapq.heappush(end_times, intervals[0][1])
361
362    for start, end in intervals[1:]:
363        # 가장 빨리 끝나는 회의가 현재 회의 시작 전에 끝나면 재사용
364        if end_times[0] <= start:
365            heapq.heappop(end_times)
366        heapq.heappush(end_times, end)
367
368    return len(end_times)
369
370
371# =============================================================================
372# 테스트
373# =============================================================================
374
375def main():
376    print("=" * 60)
377    print("힙과 우선순위 큐 (Heap & Priority Queue) 예제")
378    print("=" * 60)
379
380    # 1. 최소 힙
381    print("\n[1] 최소 힙 (직접 구현)")
382    min_heap = MinHeap()
383    for val in [5, 3, 8, 1, 2, 7]:
384        min_heap.push(val)
385    print(f"    삽입: [5, 3, 8, 1, 2, 7]")
386    print(f"    힙 배열: {min_heap.heap}")
387    print(f"    pop 순서: ", end="")
388    result = []
389    while min_heap:
390        result.append(min_heap.pop())
391    print(result)
392
393    # 2. 최대 힙
394    print("\n[2] 최대 힙")
395    max_heap = MaxHeap()
396    for val in [5, 3, 8, 1, 2, 7]:
397        max_heap.push(val)
398    print(f"    삽입: [5, 3, 8, 1, 2, 7]")
399    print(f"    pop 순서: ", end="")
400    result = []
401    while max_heap:
402        result.append(max_heap.pop())
403    print(result)
404
405    # 3. Heapify
406    print("\n[3] Heapify (배열 → 힙)")
407    arr = [9, 5, 6, 2, 3]
408    heap = MinHeap()
409    heap.heapify(arr)
410    print(f"    원본: {arr}")
411    print(f"    힙: {heap.heap}")
412
413    # 4. 힙 정렬
414    print("\n[4] 힙 정렬")
415    arr = [64, 34, 25, 12, 22, 11, 90]
416    sorted_arr = heap_sort(arr)
417    print(f"    원본: {arr}")
418    print(f"    정렬: {sorted_arr}")
419
420    # 5. K번째 요소
421    print("\n[5] K번째 요소")
422    arr = [3, 2, 1, 5, 6, 4]
423    print(f"    배열: {arr}")
424    print(f"    2번째로 큰 수: {kth_largest(arr, 2)}")
425    print(f"    2번째로 작은 수: {kth_smallest(arr, 2)}")
426
427    # 6. 중앙값 스트림
428    print("\n[6] 중앙값 스트림")
429    mf = MedianFinder()
430    stream = [2, 3, 4]
431    print(f"    스트림: {stream}")
432    for num in stream:
433        mf.add_num(num)
434        print(f"    {num} 추가 후 중앙값: {mf.find_median()}")
435
436    # 7. K개 정렬 리스트 병합
437    print("\n[7] K개 정렬 리스트 병합")
438    lists = [[1, 4, 5], [1, 3, 4], [2, 6]]
439    merged = merge_k_sorted_lists(lists)
440    print(f"    입력: {lists}")
441    print(f"    병합: {merged}")
442
443    # 8. Top K 빈도
444    print("\n[8] Top K 빈도 요소")
445    nums = [1, 1, 1, 2, 2, 3]
446    k = 2
447    result = top_k_frequent(nums, k)
448    print(f"    배열: {nums}, k={k}")
449    print(f"    결과: {result}")
450
451    # 9. 가장 가까운 K개 점
452    print("\n[9] 원점에서 가장 가까운 K개 점")
453    points = [(1, 3), (-2, 2), (5, 8), (0, 1)]
454    k = 2
455    closest = k_closest_points(points, k)
456    print(f"    점들: {points}, k={k}")
457    print(f"    결과: {closest}")
458
459    # 10. 회의실
460    print("\n[10] 최소 회의실 수")
461    meetings = [(0, 30), (5, 10), (15, 20)]
462    rooms = min_meeting_rooms(meetings)
463    print(f"    회의: {meetings}")
464    print(f"    필요한 회의실: {rooms}개")
465
466    print("\n" + "=" * 60)
467
468
469if __name__ == "__main__":
470    main()