AI刷题：小R的数组最小化问题 | 豆包MarsCode AI 刷题

题解

这道题目要求我们实现一个函数 solution，它接受三个参数：一个整数列表 a，一个整数 k，和一个整数 x。目标是在列表 a 中进行 k 次操作，每次从当前最大值中减去 x，然后重新插入到列表中。最终返回操作后的最大值。

解题思路

1. 使用最大堆：为了高效地找到和更新当前的最大值，我们可以使用 Python 的 heapq 模块来模拟一个最大堆。由于 heapq 默认是最小堆，我们需要将数组元素取负数来模拟最大堆。
2. 初始化最大堆：将输入数组中的每个元素取负数，然后使用 heapq.heapify 将其转换为堆。
3. 执行 k 次操作：
   • 每次操作中，首先取出当前最大值（取负后是最小值，因此用 heapq.heappop）。
   • 将该最大值减去 x，然后将结果重新插入堆中（取负后放回）。
4. 返回结果：操作完成后，堆顶元素即为操作后的最大值（取负后变为最小值）。

代码实现

import heapq

def solution(a: list, k: int, x: int) -> int:
    # 将数组元素取负数，以模拟最大堆
    max_heap = [-num for num in a]
    heapq.heapify(max_heap)
    
    # 进行 k 次操作
    for _ in range(k):
        # 取出当前最大值（取负后是最小值，因此用 heapq.heappop）
        max_val = -heapq.heappop(max_heap)
        # 减少 x 后再放回堆中（取负后放回）
        heapq.heappush(max_heap, -(max_val - x))
    
    # 堆顶元素即为操作后的最大值（取负后变为最小值）
    return -max_heap[0]

if __name__ == '__main__':
    print(solution(a = [7, 2, 1], k = 3, x = 2) == 2)  # True
    print(solution(a = [10, 5, 8], k = 5, x = 3) == 4)  # True
    print(solution(a = [9, 4, 7], k = 4, x = 1) == 6)  # True

测试用例解释

1. 测试用例 1：

   • 初始数组为 [7, 2, 1]。
   • k = 3，x = 2。
   • 操作过程：
     • 第一次操作：最大值 7，减去 2，得到 5，新的数组为 [5, 2, 1]。
     • 第二次操作：最大值 5，减去 2，得到 3，新的数组为 [3, 2, 1]。
     • 第三次操作：最大值 3，减去 2，得到 1，新的数组为 [1, 2, 1]。
   • 最终最大值为 2。

2. 测试用例 2：

   • 初始数组为 [10, 5, 8]。
   • k = 5，x = 3。
   • 操作过程：
     • 第一次操作：最大值 10，减去 3，得到 7，新的数组为 [7, 5, 8]。
     • 第二次操作：最大值 8，减去 3，得到 5，新的数组为 [7, 5, 5]。
     • 第三次操作：最大值 7，减去 3，得到 4，新的数组为 [4, 5, 5]。
     • 第四次操作：最大值 5，减去 3，得到 2，新的数组为 [2, 5, 4]。
     • 第五次操作：最大值 5，减去 3，得到 2，新的数组为 [2, 2, 4]。
   • 最终最大值为 4。

3. 测试用例 3：

   • 初始数组为 [9, 4, 7]。
   • k = 4，x = 1。
   • 操作过程：
     • 第一次操作：最大值 9，减去 1，得到 8，新的数组为 [8, 4, 7]。
     • 第二次操作：最大值 8，减去 1，得到 7，新的数组为 [7, 4, 7]。
     • 第三次操作：最大值 7，减去 1，得到 6，新的数组为 [6, 4, 7]。
     • 第四次操作：最大值 7，减去 1，得到 6，新的数组为 [6, 4, 6]。
   • 最终最大值为 6。