Title (medium)

You are visiting a farm where a row of fruit trees are planted from left to right. These trees are represented by an integer array of fruits, where fruits[i] is the fruit species on the ith tree.

You want to collect as many fruits as possible. However, the owner of the farm has set some strict rules. You must pick fruit according to the requirements:

You only have two baskets, and each basket can only hold a single type of fruit. There is no limit to the total amount of fruit each basket can hold.
You can choose any tree to start picking. You must pick exactly one fruit from each tree (including the tree that starts picking). The fruit picked should conform to the type of fruit in the basket. Each time you pick, you will move to the next tree to the right and continue picking.
Once you go to a tree and the fruit doesn't match the type of fruit in the basket, you must stop picking.
Give you an integer array fruits and return the maximum number of fruits you can collect.

Example 1:

Input: fruits = [1,2,1]
Output: 3
Explanation: all 3 trees can be picked.
Example 2:

Input: fruits = [0,1,2,2]
Output: 3
Explanation: you can pick [1,2,2] these three trees.
If you start picking from the first tree, you can only pick [0,1] these two trees.
Example 3:

Input: fruits = [1,2,3,2,2]
Output: 4
Explanation: you can pick [2,3,2,2] these four trees.
If you start picking from the first tree, you can only pick [1,2] these two trees.
Example 4:

Input: Fruits = [3,3,3,1,2,1,2,3,4]
Output: 5
Explanation: you can pick [1,2,1,1,2] these five trees.

Tips:

1 <= fruits.length <= 105
0 <= fruits[i] < fruits.length

Problem solving ideas

Depending on the amount of data, O(N^2) must timeout. Consider dp. dp[i] represents the number of fruits taken forward according to the rules up to position i. the maximum value in the dp array is the answer.

First find out the third fruit and deal with it from scratch until you encounter the third fruit; Record the position of the third fruit_ pos;

Then, starting from the third fruit, it was recorded that the near fruit was the third fruit, and the far far fruit was the third fruit_ Fruit at pos-1 position;
First calculate dp[third_pos], that is, take far from the starting point of continuous far fruit_ Num far fruit, and then take 1 third_ The third fruit on POS,
dp[third_pos] = far_num + 1;

Then from third_pos+1 starts to traverse backward. If the fruit is near, you can take it continuously, dp[i]=dp[i-1]+1, and near at the same time_ Num also adds one;
If the fruit is far, you can also take it together, dp[i]=dp[i-1]+1, but exchange near and far, because the adjacent fruit has become far fruit, and the occurrence times are near_num should become 1;
If there are fruits other than near and far, only the fruit and near fruit can be taken, that is, the dp value taken forward according to the rule at this position will become dp[i]=near_num + 1； And update far to near fruit, near fruit is the current fruit, and the occurrence times are near_num also becomes 1.

Finally, the maximum value in the dp array is the answer.

In addition, the dp update is only related to the previous location and can be optimized with only one int.

code

class Solution {
public:
    int totalFruit(vector<int>& fruits) {
        int n = fruits.size();
        if(n < 3) return n;
        int dp = 0;
        int ans = 0;
        int near = -1;
        int near_num = 0;
        int far = fruits[0];
        int far_num = 0;
        int third_pos = 0;
        for(int i = 1; i < n; i++) {
            if(near == -1 && fruits[i] != far) {
                near = fruits[i];
            }
            else if(fruits[i] != far && fruits[i] != near) {
                third_pos = i;
                near = fruits[i];
                near_num = 1;
                break;
            }
        }
        if(third_pos == 0) return n;    //There are only two kinds of fruit in total. You can take both
        ans = third_pos;    //Update the answer to start from scratch and take the first two fruits
        far = fruits[third_pos-1];
        for(int i = third_pos-1; i >= 0; i--) {
            if(fruits[i] == far) far_num++;
            else break;
        }
        dp = far_num + 1;    //Calculate the third fruit and the fruit before it
        for(int i = third_pos + 1; i < n; i++) {
            if(fruits[i] == near) {
                dp++;
                near_num++;
            }
            else if(fruits[i] == far) {
                dp++;
                swap(near, far);
                near_num = 1;
            }
            else {
                far = near;
                near = fruits[i];
                dp = near_num + 1;
                near_num = 1;
            }
            ans = max(ans, dp);
        }
        return ans;
    }
};