Description

Submission

class Solution {
    template<typename Iterator>
    inline int rangeSum(Iterator& it1, Iterator& it2) {
        int ret = 0;
        for(auto it = it1; it != it2; ++it) {
            ret += it->second;
        }
        return ret;
    }
    map<int, int> Map;
    
    inline int getFromMap(int index) {
        auto it = Map.find(index);
        if(it == Map.end()) return 0;
        return it->second;
    }
public:
    int maxTotalFruits(vector<vector<int>>& fruits, int startPos, int k) {
        for(auto& fruit: fruits) {
            Map[fruit[0]] = fruit[1];
        }
        
        // 2. go right all the way, then step by step go left
        int endPos = startPos + k;
        auto it1 = Map.lower_bound(startPos);
        auto it2 = Map.upper_bound(endPos);
        int total = rangeSum(it1, it2);
        int ret = total;    // max value
        for(int i = 0; i < k; ++i) {
            // cur to remove
            int cur = endPos - i;
            total -= getFromMap(cur);
            int nLeft = k - (cur - startPos);
            int l1 = cur - nLeft - 1;
            int l2 = cur - nLeft - 2;
            if(l1 < startPos) total += getFromMap(l1);
            if(l2 < startPos) total += getFromMap(l2);
            ret = max(ret, total);
        }
        
        // 1. go left all the way, then step by step go right
        endPos = startPos - k;
        it1 = Map.lower_bound(endPos);
        it2 = Map.upper_bound(startPos);
        total = rangeSum(it1, it2);
        ret = max(total, ret);
        for(int i = 0; i < k; ++i) {
            int cur = endPos + i; // cur to remove
            total -= getFromMap(cur);
            int nRight = k - (startPos - cur);
            int r1 = cur + nRight + 1;
            int r2 = cur + nRight + 2;
            if(r1 > startPos) total += getFromMap(r1);
            if(r2 > startPos) total += getFromMap(r2);
            ret = max(ret, total);
        }
        return ret;
    }
};

// find a left-most or right-most point such that the result is maximized
// 2 possible ways:
// 1. go left all the way, then step by step go right
// 2. go right all the way, then step by step go left