摘水果

标签: 数组 二分查找 前缀和 滑动窗口

难度: Hard

在一个无限的 x 坐标轴上,有许多水果分布在其中某些位置。给你一个二维整数数组 fruits ,其中 fruits[i] = [positioni, amounti] 表示共有 amounti 个水果放置在 positioni 上。fruits 已经按 positioni 升序排列 ,每个 positioni 互不相同

另给你两个整数 startPosk 。最初,你位于 startPos 。从任何位置,你可以选择 向左或者向右 走。在 x 轴上每移动 一个单位 ,就记作 一步 。你总共可以走 最多 k 步。你每达到一个位置,都会摘掉全部的水果,水果也将从该位置消失(不会再生)。

返回你可以摘到水果的 最大总数

示例 1:

输入:fruits = [[2,8],[6,3],[8,6]], startPos = 5, k = 4
输出:9
解释:
最佳路线为:
- 向右移动到位置 6 ,摘到 3 个水果
- 向右移动到位置 8 ,摘到 6 个水果
移动 3 步,共摘到 3 + 6 = 9 个水果

示例 2:

输入:fruits = [[0,9],[4,1],[5,7],[6,2],[7,4],[10,9]], startPos = 5, k = 4
输出:14
解释:
可以移动最多 k = 4 步,所以无法到达位置 0 和位置 10 。
最佳路线为:
- 在初始位置 5 ,摘到 7 个水果
- 向左移动到位置 4 ,摘到 1 个水果
- 向右移动到位置 6 ,摘到 2 个水果
- 向右移动到位置 7 ,摘到 4 个水果
移动 1 + 3 = 4 步,共摘到 7 + 1 + 2 + 4 = 14 个水果

示例 3:

输入:fruits = [[0,3],[6,4],[8,5]], startPos = 3, k = 2
输出:0
解释:
最多可以移动 k = 2 步,无法到达任一有水果的地方

提示:

  • 1 <= fruits.length <= 105
  • fruits[i].length == 2
  • 0 <= startPos, positioni <= 2 * 105
  • 对于任意 i > 0positioni-1 < positioni 均成立(下标从 0 开始计数)
  • 1 <= amounti <= 104
  • 0 <= k <= 2 * 105

Submission

运行时间: 124 ms

内存: 49.3 MB

class Solution:
    def maxTotalFruits(self, fruits: List[List[int]], startPos: int, k: int) -> int:
        left = bisect_left(fruits, [startPos-k])  #向左能走到的最远的地方
        right = bisect_left(fruits, [startPos+1])  # 起始位置之后的有水果的位置
        ans = s = sum(fruits[i][1] for i in range(left, right))
        # 滑动窗口
        n = len(fruits)
        while right < n and fruits[right][0] - startPos <= k:
            s += fruits[right][1]
            # 注意这里有两种方式:先左后右和先右后左
            while startPos + fruits[right][0] - 2 * fruits[left][0] > k and \
                2 * fruits[right][0] - startPos - fruits[left][0] > k:
                s -= fruits[left][1]
                left += 1
            ans = max(ans, s)
            right += 1
        return ans

Explain

此题解采用滑动窗口与二分搜索的结合方法,以最大化在给定步数内摘到的水果总数。首先,使用两次二分搜索确定起始位置左右的水果范围,即可能达到的左边界和右边界。随后,通过滑动窗口遍历这个范围,尝试不同的路径组合(先左后右或者先右后左),计算在不超过k步的条件下可以摘到的最大水果数。每次移动窗口时,更新当前窗口内的水果总数,并与之前的最大值进行比较,以保留可能的最大收获。

时间复杂度: O(n)

空间复杂度: O(1)

# Solution class with maxTotalFruits method
class Solution:
    def maxTotalFruits(self, fruits: List[List[int]], startPos: int, k: int) -> int:
        # Use binary search to find the furthest position to the left
        left = bisect_left(fruits, [startPos-k])
        # Find the first fruit position right of the startPos
        right = bisect_left(fruits, [startPos+1])
        # Initialize the sum of fruits within the window
        ans = s = sum(fruits[i][1] for i in range(left, right))
        # Total number of fruit positions
        n = len(fruits)
        # Expand the window towards the right within the allowed steps
        while right < n and fruits[right][0] - startPos <= k:
            s += fruits[right][1]
            # Adjust the window by moving the left side, ensuring total steps do not exceed k
            while startPos + fruits[right][0] - 2 * fruits[left][0] > k and \
                2 * fruits[right][0] - startPos - fruits[left][0] > k:
                s -= fruits[left][1]
                left += 1
            # Update the maximum fruits collected
            ans = max(ans, s)
            right += 1
        return ans

Explore

在使用二分搜索确定起始位置左右的水果范围时,并没有直接确保所有选定范围内的水果都可以在k步内被摘到。二分搜索是用来快速定位到与起始位置距离可能在k步以内的水果的大致区域。具体是否能摘到,还需要通过滑动窗口的调整和计算来验证。通过调整窗口的左右边界,确保窗口内的水果位置总和与起始位置的距离不超过k步。

确实,示例代码中主要展示了从起始位置向右扩展窗口的情况。要处理另一种方向即先向左再向右,可以通过初始化窗口时向左扩展,然后逐渐减少左边界的水果,同时增加右边界的水果,保持总步数不超过k。这种情况可以通过镜像地应用相同的逻辑来实现,即先设定一个向左的窗口,然后在必要时向右扩展。

在移动窗口左边界时,每次移动前都会对当前窗口内的水果总数进行比较和记录,确保记录下最大值。只有在确认右移左边界后,即使总步数减少,也能使得窗口内新的水果组合(通过增加右边界)可能带来更大的收获时,才会进行移动。这通过不断更新最大水果数的方式确保不会错过可能的最大收获。

在代码中使用`startPos+1`来找到右侧水果的起始位置是为了确保从起始位置的右侧开始计数水果,而不包括位于`startPos`位置上的水果。这是因为算法可能需要分别计算从起始点向左和向右两个方向上的水果收获,故需要明确界定起始点右侧的第一个位置。如果使用`startPos`,则可能会重复计算起始点位置上的水果。