2558. 从数量最多的堆取走礼物

🎯 问题描述

描述：
给你一个整数数组 gifts，表示各堆礼物的数量。每一秒，你需要执行以下操作：

  1. 选择礼物数量最多的那一堆（如果有多堆数量相同且最多，任选其一）。
  2. 将该堆的礼物数量减少到原来的平方根（向下取整）。

返回 k 秒后剩下的礼物总数量。

说明：

  • 这是一个典型的**堆（优先队列）**应用问题。

  • 核心策略：

    • 题目要求每次都要找到当前“最大值”进行操作，这正是最大堆（大根堆）的拿手好戏。
    • 由于 Python 的 heapq 模块默认只提供最小堆，我们依然采用将数值取相反数（负数）的方法，来巧妙地模拟最大堆。
    • 每次从堆顶取出绝对值最大的负数（即最多的礼物堆），计算其平方根后，再将结果的相反数放回堆中。

  • 示例：

  • 示例 1：

1
2
3
4
5
6
7
8
输入：gifts = [25,64,9,4,100], k = 4
输出：29
解释：
第1秒：选100，变为 sqrt(100)=10，数组变为 [25,64,9,4,10]
第2秒：选64，变为 sqrt(64)=8，数组变为 [25,8,9,4,10]
第3秒：选25，变为 sqrt(25)=5，数组变为 [5,8,9,4,10]
第4秒：选10，变为 sqrt(10)=3，数组变为 [5,8,9,4,3]
剩余总和：5+8+9+4+3 = 29

💻 解题思路

思路1：最大堆（优先队列）模拟

核心逻辑

  1. 构建最大堆：
    • 遍历 gifts 数组，将所有礼物数量取相反数（-num），存入列表 nums。
    • 使用 heapq.heapify(nums) 将列表原地转换为最小堆。此时堆顶（最小的负数）对应的就是原始数据中最大的礼物堆。
  2. 模拟 k 次取礼物操作：
    • 循环 k 次。每次从堆顶弹出最小的负数 candidate。
    • 将其还原为正数并计算平方根：int((-candidate)**0.5)。
    • 将计算结果的相反数（-n）重新压入堆中，等待下一轮的比较。
  3. 计算剩余总和：
    • 操作结束后，堆中存储的是所有礼物堆最终数量的相反数。
    • 对堆中所有元素求和并取相反数（-sum(nums)），即可得到最终剩下的礼物总数。

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import heapq
from typing import List

class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        # 1. 将所有礼物数量取反，构建最小堆来模拟最大堆
        nums = [-num for num in gifts]
        heapq.heapify(nums)
        
        # 2. 模拟 k 秒的操作
        for _ in range(k):
            # 弹出当前礼物最多的堆（最小的负数）
            candidate = heapq.heappop(nums)
            # 计算平方根并向下取整，再取反放回堆中
            n = int((-candidate) ** 0.5)
            heapq.heappush(nums, -n)
        
        # 3. 堆中所有元素求和并取反，得到剩余的礼物总数
        return -sum(nums)
        
        
"""
class Solution:
    def pickGifts(self, gifts: List[int], k: int) -> int:
        heapify_max(gifts)
        for _ in range(k):
            candidate = heappop_max(gifts)
            n = int(candidate**0.5)
            heappush_max(gifts,n)
        return sum(gifts)
"""

📊 性能分析

  • 时间复杂度：O((n + k) log n)
    • 构建初始堆的时间复杂度为 O(n)。
    • 每次 heappop 和 heappush 操作的时间复杂度为 O(log n)。
    • 我们一共进行 k 次操作，因此总时间复杂度为 O((n + k) log n)。
  • 空间复杂度：O(n)
    • 需要一个额外的堆（列表）来存储 n 个礼物堆的数量。

思考

  1. 为什么要用负数模拟最大堆？

    • 原理与“最后一块石头的重量”完全一致。Python 的 heapq 只能弹出最小值。通过存入负数，原本最大的正数（如 100）变成了最小的负数（-100），正好位于堆顶，让我们每次都能 O(1) 地获取到当前最大的礼物堆。

  2. 关于平方根的计算：

    • 代码中使用了 int((-candidate) ** 0.5)。** 0.5 是计算平方根的常用方法，int() 会自动向下取整，完美符合题目“减少到原来的平方根（向下取整）”的要求。

  3. 最后求和的巧妙之处：

    • 很多类似的题目最后只需要返回堆顶元素（如前几题）。但这道题要求返回剩余总和。
    • 你的代码直接使用 -sum(nums)，省去了逐个弹出堆元素的步骤，非常简洁高效。

  4. 数据范围的考量：

    • 随着操作的进行，礼物数量会迅速减小（因为开平方根会让大数迅速变小）。这意味着即使 k 很大，堆中的数值也会很快收敛到较小的值（比如 1 的平方根还是 1），程序的运行效率会非常高。