LeetCode 1705 吃苹果的最大数目

class Solution:
    def eatenApples(self, apples: List[int], days: List[int]) -> int:
        ans = 0
        pq = []
        i = 0
        # 第一阶段还在产苹果
        while i < len(apples):
            while pq and pq[0][0] <= i:
                heappop(pq)
            # 如果树上长苹果了
            if apples[i]:
                heappush(pq, [i + days[i], apples[i]])
            if pq:
                # 当天吃掉一个苹果
                pq[0][1] -= 1
                if pq[0][1] == 0:
                    heappop(pq)
                ans += 1
            i += 1
        # 第二阶段不产苹果了
        while pq:
            while pq and pq[0][0] <= i:
                heappop(pq)
            if len(pq) == 0:
                break
            p = heappop(pq)
            # 过期时间和苹果数量的较小值，省了几步
            num = min(p[0] - i, p[1])
            ans += num
            i += num
        return ans