본문 바로가기

코딩테스트/알고리즘

예제로 배우는 알고리즘 - 힙

최소 혹은 최대값을 꺼내며 확인

 

예제 - LeetCode 215. Kth Largest Element in an Array

https://leetcode.com/problems/kth-largest-element-in-an-array/description/?envType=study-plan-v2&envId=top-interview-150

 

1. 문제 이해

  • 주어진 배열에서 K번째 큰 요소 반환
    • 중복 제거, 정렬 없이 풀이

2. 요구사항 분석

  • 입력
    • 1 <= k <= nums.length <= 100,000
    • -10,000 <= nums[i] <= 10,000
  • 출력
    • K번째로 큰 요소

 

 

3. 해결 아이디어

  • 최소 힙을 우선순위 큐를 활용
  • 힙에 값을 넣으며 크기가 k를 초과하면 해당 값 제거
  • 마지막에 가장 위에 있는 값이(가장 작은 값)이 k번째로 큰 값

 

4. 시간 복잡도 및 공간 복잡도 확인

  • 시간 복잡도
    • k크기를 n번 힙 연산 O(nlogk)
  • 공간 복잡도
    • k크기의 우선순위 큐가 필요하므로 O(k)

 

5. 구현

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // 1. 최소 힙 선언
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        
        // 2. 모든 요소 순회
        for (int num : nums) {
            // 2-1) 힙에 요소 삽입
            minHeap.offer(num);
            // 2-2)힙에 크기가 k를 초과하면 가장 작은 값 제거
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
        
        // 3. k번째로 큰 값이 힙의 top에 있음
        return minHeap.peek();
    }
}

 

예제 - LeetCode 502. IPO

https://leetcode.com/problems/ipo/description/?envType=study-plan-v2&envId=top-interview-150

 

1. 문제 이해

  • 최대 k개의 프로젝트를 선택하여 최종 자본 최대화하기
    • 각 프로젝트는 시작을 위한 필요 자본과 완료시 이익이 존재

2. 요구사항 분석

  • 입력
    • 1 <= k <= 100,000
    • 0 <= w <= 1,000,000,000
    • 1 <= profits.length == capital.length <= 100,000
    • 0 <= profits[i] <= 10,000
    • 0 <= capital[i] <= 1,000,000,000
  • 출력
    • 최대 최종 자본

 

 

3. 해결 아이디어

  • capital 기준 프로젝트 정렬
  • 현재 자본으로 수행 가능한 프로젝트들을 최대 힙(profit 기준)에 추가
  • 가장 이익이 큰 프로젝트부터 선택하여 진행

 

4. 시간 복잡도 및 공간 복잡도 확인

  • 시간 복잡도
    • 정렬(nlogn) + k번의 힙연산(klogn) O(nlogn + klogn)
  • 공간 복잡도
    • n크기의 우선순위 큐가 필요하므로 O(k)

 

5. 구현

class Solution {
    public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
        int n = profits.length;

        // 1. 프로젝트를 capital 기준으로 정렬
        List<int[]> projects = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            projects.add(new int[]{capital[i], profits[i]});
        }
        projects.sort(Comparator.comparingInt(p -> p[0]));

        // 2. max-heap으로 이익이 큰 순서대로 처리
        PriorityQueue<Integer> maxProfitHeap = new PriorityQueue<>(Collections.reverseOrder());

        int i = 0;
        while (k-- > 0) {
            // 2-1) 현재 자본 w로 가능한 프로젝트들 추가
            while (i < n && projects.get(i)[0] <= w) {
                maxProfitHeap.offer(projects.get(i)[1]);
                i++;
            }

            // 2-2) 수행 가능한 프로젝트가 없으면 중단
            if (maxProfitHeap.isEmpty()) {
                break;
            }

            // 2-3)가장 이익이 큰 프로젝트 수행
            w += maxProfitHeap.poll();
        }

        return w;
    }
}

 

예제 - LeetCode 373. Find K Pairs with Smallest Sums

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/?envType=study-plan-v2&envId=top-interview-150

 

1. 문제 이해

  • 두 개의 정렬된 배열에서 하나씩 뽑아 만든 쌍 중 합이 가장 작은 k개의 쌍 구하기
    • 각 쌍의 합이 작을수록 먼저 반환

2. 요구사항 분석

  • 입력
    • 1 <= k <= 10,000
    • 1 <= nums1.length, nums2.length <= 100,000
    • k <= nums1.length * nums2.length
    • -1,000,000,000 <= nums1[i], nums2[i] <= 1,000,000,000
  • 출력
    • 합이 가장 작은 k개의 쌍 리스트

 

 

3. 해결 아이디어

  • 브루트 포스 : 모든 쌍 조합 확인 -> O(n * m)으로 최대 10,000,000,000번 연산하므로 시간 초과
  • Min Heap + 인덱스
    • 합과 각 인덱스를 힙에 저장
    • nums1[i] + nums2[0]의 초기 쌍을 k개까지 힙에 삽입 
    • 쌍의 합이 작은 순서로 꺼내며 nums2의 인덱스를 확장하며 힙에 추가

 

4. 시간 복잡도 및 공간 복잡도 확인

  • 시간 복잡도
    •  k번의 힙연산 O(klogk)
  • 공간 복잡도
    • n크기의 우선순위 큐가 필요하므로 O(k)

 

5. 구현

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();

        // 1. 합, 각 인덱스를 Min Heap에 추가(합 기준)
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));

        // 2. nums1[i] + nums2[0] 형태의 쌍을 k개까지 삽입
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            minHeap.offer(new int[]{nums1[i] + nums2[0], i, 0});
        }

        // 3. 결과 추가 및 확장
        while (k-- > 0 && !minHeap.isEmpty()) {
            // 3-1) 쌍의 합이 작은 순서로 결과에 추가
            int[] current = minHeap.poll();
            int i = current[1], j = current[2];
            result.add(Arrays.asList(nums1[i], nums2[j]));

            // 3-2) num2를 확장한 결과를 힙에 추가
            if (j + 1 < nums2.length) {
                minHeap.offer(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
            }
        }

        return result;
    }
}

 

예제 - LeetCode 295. Find Median from Data Stream

https://leetcode.com/problems/find-median-from-data-stream/description/?envType=study-plan-v2&envId=top-interview-150

 

1. 문제 이해

  • 중앙값을 찾는 MedianFinder 객체 만들기
    • MedianFinder() : 객체 초기화
    • void addNum(int num) : 데이터 스트림에서 숫자 추가
    • double findMedian() : 현재까지의 중앙값 반환

2. 요구사항 분석

  • 입력
    • -100,000 <= num <= 100,000
    • 최대 50,000번 메소드가 호출됨
  • 출력
    • 중앙값은 double 형식으로 반환

 

 

3. 해결 아이디어

  • 최소힙과 최대힙을 활용
    • 숫자 추가 시 최대힙에 넣고, 최대힙에서 가장 큰 값을 최소 힙에 추가
    • 최대힙이 항상 최소힙보다 크기가 크도록 유지
    • 전체 요소가 홀수인 경우 최대힙의 가장 위의 값이 중앙값
    • 전체 요소가 짝수인 경우 최대힙과 최소힙의 가장 위의 값의 평균이 중앙값

 

4. 시간 복잡도 및 공간 복잡도 확인

  • 시간 복잡도
    •  addNum : 힙 연산을 하기에 O(logn)
    • findeMedian : O(1)
  • 공간 복잡도
    • n크기의 우선순위 큐가 두개 필요하므로 O(n)

 

5. 구현

class MedianFinder {
    // 최대 힙: 중앙값 이하
    private PriorityQueue<Integer> maxHeap;
    // 최소 힙: 중앙값 초과
    private PriorityQueue<Integer> minHeap;

    public MedianFinder() {
        maxHeap = new PriorityQueue<>(Collections.reverseOrder());
        minHeap = new PriorityQueue<>();
    }

    public void addNum(int num) {
        maxHeap.offer(num);
        minHeap.offer(maxHeap.poll());

        // left가 항상 right보다 크거나 같도록 유지
        if (maxHeap.size() < minHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }

    public double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.peek(); // 홀수
        } else {
            return (maxHeap.peek() + minHeap.peek()) / 2.0; // 짝수
        }
    }
}