904. Fruit Into Baskets

传送门

问题分析

本问题重点在优化从何处再次搜索tree，这里我使用lastLoc用于记录不同类型最后一次出现的位置下标。那么我们再次搜索时只用从min(lastLoc[0], lastLoc[1]) + 1 出发就可以了，这样可以从一定程度上加速搜索速度。

源码

class Solution {
public:
    int totalFruit(vector<int>& tree) {
        int sum{0}, amount{0}, typeLoc = {0}, loc;
        int type[2] = {-1, -1};
        int lastLoc[2] = {-1, -1};
        for(int i = 0; i < tree.size(); i++){
            if(typeLoc < 2){
                if(typeLoc == 1 && type[0] == tree[i]){
                    lastLoc[typeLoc] = i;
                }
                else{
                    type[typeLoc] = tree[i];
                    lastLoc[typeLoc] = i;
                    typeLoc++;
                }
                sum++;
            }
            else{
                // type in basket
                if(type[0] == tree[i] || type[1] == tree[i]) {
                    if(type[0] == tree[i])  lastLoc[0] = i;
                    if(type[1] == tree[i])  lastLoc[1] = i;
                    sum++;
                }
                // type not in basket
                else {
                    i = min(lastLoc[0], lastLoc[1]);
                    for(int i = 0; i < 2; i++){
                        type[i] = -1;
                        lastLoc[i] = -1;
                    }
                    amount = max(sum, amount);
                    typeLoc = 0;
                    sum = 0;
                }
            }
        }
        return max(amount, sum);            
    }
};