LeetCode/Top Interview 150

502. IPO

hyunkookim 2024. 12. 22. 18:48

502. IPO

 

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)으로 풀어야 함!!