[前缀和+二分] 2106.摘苹果

题目

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

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

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

思路分析

• 一开始考虑两边分别求前缀和，但是超时，而且还要维护两个前缀和区间和下标
• 边界处理要考虑开区间和闭区间，lower_bound在区间内无需改动，upper_bound需要-1,才能得到第一个小于等于的下标
• 注意前缀和的sum[pr] - sum[pl]+fruits[pl][1]

代码如下

 

class Solution {
public:
    int maxTotalFruits(vector<vector<int>>& fruits, int startPos, int k) {
        int n = fruits.size();

        // 前缀和，用于求区间 [l, r] 共有多少水果
        vector<int> sum;
        for (int i = 0; i < n; ++ i){
            if(i==0) sum.push_back(fruits[0][1]);
            else sum.push_back(fruits[i][1] + sum[i-1]);
        }
        vector<int> pos;
        for (int i = 0; i < n; ++ i)
            pos.push_back(fruits[i][0]);

        // 枚举可能抵达的区间
        int ans = 0;
        for (int x = k; x >= 0; -- x)
        {
            int y = (k - x) / 2;
            int l, r;
            // x + 2y = k
            l = startPos - x, r = startPos + y;
            auto pl = lower_bound(pos.begin(), pos.end(), l) - pos.begin(), pr = upper_bound(pos.begin(), pos.end(), r) - pos.begin()-1;
            if(pl!=pos.size()&&pr!=-1) ans = max(ans, sum[pr] - sum[pl]+fruits[pl][1]);
            // 2y + x = k
            l = startPos - y, r = startPos + x;
            pl = lower_bound(pos.begin(), pos.end(), l) - pos.begin(), pr = upper_bound(pos.begin(), pos.end(), r) - pos.begin()-1;
            if(pl!=pos.size()&&pr!=-1) ans = max(ans, sum[pr] - sum[pl]+fruits[pl][1]);
        }

        return ans;
    }
};