LeetCode/NeetCode

[Two Heap] 502. IPO ★★★★★

hyunkookim 2024. 12. 22. 18:48

502. IPO

 

이 문제는 Greedy + Heap을 활용해서 해결할 수 있는 IPO 자본 최대화 문제입니다.
주어진 프로젝트 중 최대 k개를 골라 이익을 최대화하는 것이 목표입니다.


✅ 문제 요약

  • 각 프로젝트는 다음 정보를 가짐:
    • profits[i]: 이익
    • capital[i]: 최소 시작 자본
  • 시작 자본: w
  • 프로젝트를 최대 k개까지 선택할 수 있음
  • 조건:
    • 현재 자본 w 이상을 요구하는 프로젝트만 수행 가능
    • 프로젝트를 수행하면 그 이익만큼 자본 증가

✅ 핵심 아이디어

🎯 항상 지금 당장 수행 가능한 프로젝트 중에서 이익이 가장 큰 것을 선택

✔ 자료구조 선택

  • 모든 프로젝트를 (capital, profit) 형태로 정렬하여 준비
  • 매 순간, 현재 자본 w로 수행 가능한 프로젝트를 **최대 힙(max heap)**에 저장
  • 힙에서 가장 큰 profit을 뽑아서 수행 → 자본 증가 → 반복

✅ 알고리즘 흐름

  1. (capital, profit) 튜플로 묶어 정렬
  2. 현재 자본 w로 가능한 프로젝트들을 max-heap에 push
  3. heap이 비어 있지 않다면, 가장 profit 큰 것 pop → w += profit
  4. 최대 k번까지 반복

https://youtu.be/1IUzNJ6TPEM

 

Two Heap: Greedy + Heap 

import heapq
from typing import List

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        maxProfit = []  # 최대 힙: 지금 가능한 프로젝트 중 profit이 가장 큰 것 저장 (음수로 저장해서 최대 힙처럼 사용)

        # 모든 프로젝트를 (필요 자본, 이익) 형태로 묶어 최소 힙으로 구성
        # 자본(capital)이 적은 순으로 프로젝트를 쉽게 꺼내기 위함
        minCapital = [(c, p) for c, p in zip(capital, profits)]
        heapq.heapify(minCapital)  # 최소 힙 구성 (capital 오름차순)

        # 최대 k개의 프로젝트를 선택
        for _ in range(k):
            # 현재 자본 w로 시작 가능한 모든 프로젝트를 최대 힙에 넣음
            while minCapital and minCapital[0][0] <= w:
                c, p = heapq.heappop(minCapital)    # 자본이 적은 프로젝트 꺼냄
                heapq.heappush(maxProfit, -p)       # profit은 음수로 넣어 최대 힙처럼 사용

            # 수행 가능한 프로젝트가 없으면 종료
            if not maxProfit:
                break

            # 가장 이익이 큰 프로젝트 수행
            w += -heapq.heappop(maxProfit)  # 음수로 저장되어 있으니 다시 양수로 바꿔 더함

        return w  # k개 또는 가능한 모든 프로젝트 수행 후 자본 반환

 

k = 2
w = 0
profits = [1, 2, 3]
capital = [0, 1, 1]

# 1단계: w=0 → 가능한 프로젝트 = [0,1]
# maxProfit = [1]
# 선택: +1 → w = 1

# 2단계: w=1 → 가능한 프로젝트 = [1,2]
# maxProfit = [3,2]
# 선택: +3 → w = 4

📌 핵심 흐름 요약

  1. (capital[i], profits[i])를 최소 힙으로 관리해서 현재 자본으로 가능한 프로젝트만 꺼냄
  2. 가능한 프로젝트들은 최대 힙에 저장하여 가장 이익 큰 것부터 선택
  3. 최대 k번 반복하면서 자본을 늘려감

✅ 시간/공간 복잡도

  • 시간 복잡도:
    • O(n) for heapify
    • O(k log n) for up to k heap operations
    • 전체: O(n + k log n)
  • 공간 복잡도:
    • O(n) for heaps

 


 

 

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        # 1번 할때마다 k-- 감소
        # w 는 현재 자본
        # capital[]은 profits[] 얻기 위한 최소 요구 자본
        # 그래서, 매 라운드 마다 현재 자본 w 내에서 얻을수있는 최대 profit 선택하고
        # 그 profit을 w에 더하면 됨
        # 그리고 w return

         # 이미 선택된 프로젝트의 인덱스를 추적
        visited_idx = set()
        while k>0:
            # 현재 자본으로 실행 가능한 프로젝트의 인덱스를 필터링
            valid_capital_idx = [i for i, c in enumerate(capital) if c <= w and i not in visited_idx]

            if not valid_capital_idx: # 비었으면
                break

            # 실행 가능한 프로젝트 중 최대 profit을 찾음
            max_idx = valid_capital_idx[0]
            max_profits = profits[max_idx]
            
            for i in valid_capital_idx:
                if (max_profits < profits[i]):  # 최대 profit 갱신
                    max_profits = profits[i]
                    max_idx = i

            # 최대 profit을 선택하여 w 갱신
            w += max_profits        
            visited_idx.add(max_idx) # 선택된 프로젝트를 기록
            k-=1

        return w

 

Time Limit Exceeded 발생!!

 

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        # k: 최대 실행 횟수
        # w: 초기 자본
        # profits: 각 프로젝트의 수익
        # capital: 각 프로젝트의 최소 요구 자본

        while k > 0:
            # 현재 자본으로 실행 가능한 프로젝트의 인덱스 필터링
            valid_capital_idx = [i for i, c in enumerate(capital) if c <= w]

            if not valid_capital_idx:  # 실행 가능한 프로젝트가 없으면 종료
                break

            # 실행 가능한 프로젝트 중 최대 profit 찾기
            max_idx = max(valid_capital_idx, key=lambda i: profits[i])

            # 최대 profit을 선택하여 자본(w) 업데이트
            w += profits[max_idx]

            # 선택된 프로젝트 제거 (visited 관리를 대체)
            capital[max_idx] = float('inf')  # 선택된 프로젝트는 다시 고려되지 않도록 설정

            k -= 1  # 실행 횟수 감소

        return w

 

좀더 빨라지긴 했으나, 이것도 Time Limit Exceeded 발생!!

 

class Solution:
    def findMaximizedCapital(self, k: int, w: int, profits: List[int], capital: List[int]) -> int:
        """
        Python에서는 heapq가 기본적으로 최소 힙만 지원하기 때문에, 
        값을 음수로 저장해서 최대 힙처럼 동작하게 구현
        """

        # 프로젝트를 (capital, profit) 형태로 묶고, capital 기준으로 정렬
        projects = sorted(zip(capital, profits), key=lambda x: x[0])
        n = len(projects)

        idx = 0
        max_heap = [] # 최대 힙 (profits를 음수로 저장)
        
        for _ in range(k):
            # 현재 자본으로 실행 가능한 프로젝트를 힙에 추가
            while idx < n and projects[idx][0] <= w:
                heapq.heappush(max_heap, -projects[idx][1]) # 최대 힙처럼 동작하도록 음수로 저장
                idx += 1

            if not max_heap: # 실행가능한 프로젝트가 없으면
                break # 종료

            # 최대 profit 선택: 힙에서 pop
            w += -heapq.heappop(max_heap)

        return w

 

힙(heap)으로 풀어야 함!!