LeetCode刷题日记（2558. 从数量最多的堆取走礼物）

"""
你一个整数数组 gifts ，表示各堆礼物的数量。每一秒，你需要执行以下操作：
选择礼物数量最多的那一堆。
如果不止一堆都符合礼物数量最多，从中选择任一堆即可。
选中的那一堆留下平方根数量的礼物（向下取整），取走其他的礼物。
返回在 k 秒后剩下的礼物数量。
"""
from typing import List
import heapq


class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        for i in range(k):
            gifts.sort(reverse=True)
            # print(gifts)
            gifts[0] = int(gifts[0] ** 0.5)

        return sum(gifts)

    def pickGifts2(self, gifts: List[int], k: int) -> int:
        gifts = [-gift for gift in gifts]

        heapq.heapify(gifts)
        for i in range(k):
            # x = heapq.heappop(gifts)
            # x = -int((-x) ** 0.5)
            # heapq.heappush(gifts,x)
            heapq.heappush(gifts,-int((-(heapq.heappop(gifts)))**0.5))
        return -sum(gifts)