「前缀和」的经典运用——解题 1744：你能在最喜欢那日吃到最喜欢的糖果吗？

2023-11-09 01:47:07

题目

在 LeetCode 上，有这样一道题：

给你一个下标从 0 开始的正整数数组 candiesCount ，其中 candiesCount[i] 表示你拥有的第 i 颗糖果的类型数量。同时给你一个整数类型数组 queries ，其中 queries[i] 是 ZZZ 对糖果类型的询问。对于每个查询，你都要回答有多少天你至少有 queries[i] 颗这一天所询问类型的糖果。由于所有的糖果都是按日期给出的，所以你不能多次吃同一颗糖果。

注意，你拥有的糖果类型种类不超过 candiesCount.length。

输入格式

candiesCount：一个正整数数组，其中 candiesCount[i] 表示你拥有的第 i 颗糖果的类型数量。
queries：一个整数类型数组，其中 queries[i] 是 ZZZ 对糖果类型的询问。

输出格式

一个整数数组 answer，其中 answer[i] 表示对于每个查询 queries[i]，你至少有 queries[i] 颗这一天所询问类型的糖果的天数。

示例

示例 1：

输入：candiesCount = [7, 4, 5, 3, 8], queries = [2, 4, 6]
输出：[5, 4, 0]
解释：
- 对于 queries[0] = 2，有 5 天你至少有 2 颗糖果类型，分别是第 1 天、第 2 天、第 3 天、第 4 天和第 5 天。
- 对于 queries[1] = 4，有 4 天你至少有 4 颗糖果类型，分别是第 1 天、第 2 天、第 3 天和第 4 天。
- 对于 queries[2] = 6，没有一天你至少有 6 颗糖果类型，所以答案为 0。

示例 2：

输入：candiesCount = [5, 2, 6, 4, 1], queries = [3, 2, 5, 4, 1]
输出：[4, 5, 0, 1, 5]
解释：
- 对于 queries[0] = 3，有 4 天你至少有 3 颗糖果类型，分别是第 1 天、第 2 天、第 3 天和第 4 天。
- 对于 queries[1] = 2，有 5 天你至少有 2 颗糖果类型，分别是第 1 天、第 2 天、第 3 天、第 4 天和第 5 天。
- 对于 queries[2] = 5，没有一天你至少有 5 颗糖果类型，所以答案为 0。
- 对于 queries[3] = 4，有 1 天你至少有 4 颗糖果类型，就是第 4 天。
- 对于 queries[4] = 1，有 5 天你至少有 1 颗糖果类型，分别是第 1 天、第 2 天、第 3 天、第 4 天和第 5 天。

算法思路

这道题目的核心是求出每一天糖果的总数量，然后对于每个查询，找到第一个总数量大于或等于查询数量的日子。我们可以使用「前缀和」来快速计算每一天糖果的总数量，然后使用二分查找来找到第一个总数量大于或等于查询数量的日子。

代码实现

def count_candies(candiesCount, queries):
  """
  计算你至少有 queries[i] 颗这一天所询问类型的糖果的天数。

  Args:
    candiesCount: 一个正整数数组，其中 candiesCount[i] 表示你拥有的第 i 颗糖果的类型数量。
    queries: 一个整数类型数组，其中 queries[i] 是 ZZZ 对糖果类型的询问。

  Returns:
    一个整数数组 answer，其中 answer[i] 表示对于每个查询 queries[i]，你至少有 queries[i] 颗这一天所询问类型的糖果的天数。
  """

  # 计算前缀和
  prefix_sum = [0] * len(candiesCount)
  prefix_sum[0] = candiesCount[0]
  for i in range(1, len(candiesCount)):
    prefix_sum[i] = prefix_sum[i - 1] + candiesCount[i]

  # 存储答案
  answer = []

  # 遍历每个查询
  for query in queries:
    # 使用二分查找找到第一个总数量大于或等于查询数量的日子
    left, right = 0, len(candiesCount) - 1
    while left <= right:
      mid = (left + right) // 2
      if prefix_sum[mid] >= query:
        right = mid - 1
      else:
        left = mid + 1

    # 将答案添加到 answer 数组中
    answer.append(left)

  return answer


# 测试代码
candiesCount = [7, 4, 5, 3, 8]
queries = [2, 4, 6]
result = count_candies(candiesCount, queries)
print(result)  # [5, 4, 0]

candiesCount = [5, 2, 6, 4, 1]
queries = [3, 2, 5, 4, 1]
result = count_candies(candiesCount, queries)
print(result)  # [4, 5, 0, 1, 5]

复杂度分析

时间复杂度：对于每个查询，我们需要使用二分查找来找到第一个总数量大于或等于查询数量的日子，时间复杂度为 O(log n)，其中 n 是糖果数量。因此，总的时间复杂度为 O(m log n)，其中 m 是查询的数量。
空间复杂度：我们使用了一个前缀和数组来存储糖果数量的累积和，空间复杂度为 O(n)。

总结

在这篇文章中，我们详细讲解了「前缀和」在解决 LeetCode 题目 1744. 你能在你最喜欢的那天吃到你最喜欢的糖果吗？中的巧妙运用。我们使用「前缀和」快速计算每一天糖果的总数量，然后使用二分查找找到第一个总数量大于或等于查询数量的日子。这种方法大大提高了算法的效率，使我们能够在 O(m log n) 的时间内解决这个问题。希望这篇文章对您理解「前缀和」的应用有所帮助。