AI刷题 2.徒步旅行中的补给问题 题解 | 豆包MarsCode AI刷题

问题描述

小R正在计划一次从地点A到地点B的徒步旅行，总路程需要 N 天。为了在旅途中保持充足的能量，小R每天必须消耗1份食物。幸运的是，小R在路途中每天都会经过一个补给站，可以购买食物进行补充。然而，每个补给站的食物每份的价格可能不同，并且小R最多只能同时携带 K 份食物。

现在，小R希望在保证每天都有食物的前提下，以最小的花费完成这次徒步旅行。你能帮助小R计算出最低的花费是多少吗？

---

测试样例

样例1：

输入：n = 5 ,k = 2 ,data = [1, 2, 3, 3, 2]
输出：9

样例2：

输入：n = 6 ,k = 3 ,data = [4, 1, 5, 2, 1, 3]
输出：9

样例3：

输入：n = 4 ,k = 1 ,data = [3, 2, 4, 1]
输出：10

题解

代码

import heapq

def solution(n, k, data):
    # 初始化最小堆和总花费
    heap = []
    total_cost = 0
    price_map = {}

    # 遍历每一天
    for day in range(n):
        # 将当前补给站的价格加入最小堆
        heapq.heappush(heap, data[day])
        price_map[data[day]] = price_map.get(data[day], 0) + 1

        # 如果当前天数大于等于 K，移除超出 K 天范围的补给站价格
        if day >= k:
            old_price = data[day - k]
            price_map[old_price] -= 1
            if price_map[old_price] == 0:
                del price_map[old_price]

        # 从最小堆中取出价格最低的补给站购买食物
        while heap and heap[0] not in price_map:
            heapq.heappop(heap)
        total_cost += heap[0]

    return total_cost

# 测试样例
print(solution(5, 2, [1, 2, 3, 3, 2]))  # 输出: 9
print(solution(6, 3, [4, 1, 5, 2, 1, 3]))  # 输出: 9
print(solution(4, 1, [3, 2, 4, 1]))  # 输出: 10

这段代码是为了解决小R在徒步旅行中以最低花费购买食物的问题。问题的关键在于如何有效地管理小R的食物库存，以及如何选择价格最便宜的补给站来购买食物。这里使用了最小堆（优先队列）来追踪当前可访问的价格最低的补给站。

代码解释

1. 导入库:

   • heapq 是 Python 中用于实现堆数据结构的一个模块，这里用来创建一个最小堆。

2. 定义函数:

   • solution(n, k, data): 接受三个参数：n 代表天数，k 代表最多可以携带的食物份数，data 是一个列表，表示每天补给站食物的价格。

3. 初始化:

   • heap = []: 创建一个空的最小堆。
   • total_cost = 0: 初始化总花费为0。
   • price_map = {}: 用于记录堆中每个价格出现的次数。

4. 遍历每一天:

   • 对于每一天，将当天补给站的价格加入到最小堆中，并更新价格映射表。
   • 如果当前天数大于等于 k，则从价格映射表中移除超出 k 天范围的补给站价格，并调整堆。
   • 使用 while 循环确保堆顶元素是有效的（即它仍然存在于 price_map 中）。如果堆顶元素不在 price_map 中，则将其从堆中移除。
   • 从堆中取出当前价格最低的补给站，并增加总花费。

5. 返回结果:

   • 最后返回 total_cost，即完成整个旅程所需的最低花费。

具体步骤

• 维护一个最小堆 来存储当前窗口内所有补给站的价格，这样可以快速找到价格最低的那个。
• 使用字典 price_map 来跟踪堆中每个价格的数量，以便当某个价格不再有效时能够正确地从堆中移除。
• 每当超过 k 天的限制，就从 price_map 和可能的堆中移除那些不再能使用的补给站价格。
• 每天从堆中获取价格最低的补给站 并累加到总花费中。

测试样例解释

• 样例1 (n=5, k=2, data=[1, 2, 3, 3, 2]):

  • 第一天买1份，花费1；第二天买1份，花费1；第三天买1份，花费2；第四天买1份，花费2；第五天买1份，花费2。总花费9。

• 样例2 (n=6, k=3, data=[4, 1, 5, 2, 1, 3]):

  • 根据策略，最终会得到一个较低的总花费9。

• 样例3 (n=4, k=1, data=[3, 2, 4, 1]):

  • 每天都必须买当天的，因为 k=1，所以总花费10。

通过这种方法，我们可以保证每次购买都是基于当前和未来几天中最便宜的选择，从而达到最小化总花费的目的。