力扣2106-摘水果

摘要: 一个比较综合的前缀和问题

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

---

各位好，今天我们来看一个比较综合的问题，涉及到前缀和、二分和滑动窗口的综合，并且非常考验耐心。

题目

• 2106. 摘水果

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

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

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

提示：

1 <= fruits.length <= 1e5
fruits[i].length == 2
0 <= startPos, positioni <= 2 * 1e5
对于任意 i > 0 ，positioni-1 < positioni 均成立（下标从 0 开始计数）
1 <= amounti <= 1e4
0 <= k <= 2 * 1e5

示例 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 步，无法到达任一有水果的地方

题解

算法: 二分 + 前缀和 + 滑动窗口

pos[i] 为第 i 个水果的位置，ammount[i] 为第 i 个水果的数量。其中 pos 已经按照升序排列，因此最左边的位置是 pos[0]，最右边的位置为 pos[n - 1]。

从 startPos 往左走最多走到 startPos - k，往右走最多走到 startPos + k。因此我们的考虑范围是 [startPos - k, startPos + k]，超出这个范围，即使有水果也与答案无关。因此我们首先在 pos 上找到我们考虑的区间范围 [left, right]。其中：

$$left = \min\limits_{pos[i]\geq startPos-k}i \\ right = \max\limits_{pos[i]\leq startPos+k}i \\$$

有两个边界情况，如果 pos[0] >= startPos - k，那么 left = 0，如果 pos[n - 1] <= startPos + k，那么 right = n - 1。

从 startPos 出发，有两种可能取到最大值的路线，下面分别看。

• (1) 是先向左走，到达某个位置 i，满足 i >= left，且 pos[i] <= startPos，此时已经走了的距离为 startPos - pos[i]，然后再折返向右走，最多还可以走的距离为 k - startPos + pos[i]，因此最远到达 2 * pos[i] + k - startPos。

记 j1 为小于等于 startPos 的最大值；j2 为小于等于 2 * pos[i] + k - startPos 的最大值。j = max(j1, j2)，则 ammount[i..j] 的和就是一个可能的答案。

令 i 从 left 出发，向右推进，直到 j1，维护过程中的各个可能答案的最大值。

• (2) 先向右走，到达某个位置 j，满足 j <= right 且 pos[j] >= startPos，此时已经走了的距离为 pos[j] - startPos，然后再折返向左走，最多还可以走的距离为 k - pos[j] + startPos，因此最远可以走到 2 * pos[j] - k - startPos。

记 i1 为大于等于 startPos 的最小值，i2 为大于等于 2 * pos[j] - k - startPos 的最小值，i = min(i1, i2)，则 ammount[i..j] 就是一个可能的答案。

由于 pos 有序，因此前面在求 left，right, i1, j1, i2, j2 的过程都可以用二分来做。

由于我们要频繁求 ammount[i..j] 的和，可以用前缀和预处理一下 ammount 得到 sums，于是 ammount[i..j] = sums[j + 1] - sums[i]。

综上，整体的算法由前缀和、二分、滑动窗口形成。

代码 (C++)

class Solution {
public:
    int maxTotalFruits(vector<vector<int>>& fruits, int startPos, int k) {
        int n = fruits.size();
        vector<int> pos(n);
        vector<int> amount(n);
        for(int i = 0; i < n; ++i)
        {
            pos[i] = fruits[i][0];
            amount[i] = fruits[i][1];
        }

        vector<int> sums(n + 1);
        for(int i = 0; i < n; ++i)
            sums[i + 1] = sums[i] + amount[i];

        int left = lower_bound(pos.begin(), pos.end(), startPos - k) - pos.begin();
        int right = upper_bound(pos.begin(), pos.end(), startPos + k) - pos.begin() - 1;
        if(left > right)
            return 0;
        left = max(left, 0);
        right = min(right, n - 1);

        int ans = 0;

        int j1 = upper_bound(pos.begin(), pos.end(), startPos) - pos.begin() - 1;
        for(int i = left; i < n && pos[i] <= startPos; ++i)
        {
            int j2 = upper_bound(pos.begin(), pos.end(), 2 * pos[i] + k - startPos) - pos.begin() - 1;
            int j = max(j1, j2);
            ans = max(ans, sums[j + 1] - sums[i]);
        }
                
        int i1 = lower_bound(pos.begin(), pos.end(), startPos) - pos.begin();
        for(int j = right; j >= 0 && pos[j] >= startPos; --j)
        {
            int i2 = lower_bound(pos.begin(), pos.end(), 2 * pos[j] - k - startPos) - pos.begin();
            int i = min(i1, i2);
            ans = max(ans, sums[j + 1] - sums[i]);
        }

        return ans;
    }
};

---