Google

Problem

You are visiting a farm that has a single row of fruit trees arranged from left to right. The trees are represented by an integer array fruits where fruits[i] is the type of fruit the ith tree produces.

You want to collect as much fruit as possible. However, the owner has some strict rules that you must follow:

  • You only have two baskets, and each basket can only hold a single type of fruit. There is no limit on the amount of fruit each basket can hold.
  • Starting from any tree of your choice, you must pick exactly one fruit from every tree (including the start tree) while moving to the right. The picked fruits must fit in one of your baskets.
  • Once you reach a tree with fruit that cannot fit in your baskets, you must stop.

Given the integer array fruits, return the maximum number of fruits you can pick.

Example 1: Input: fruits = [1,2,1] Output: 3 Explanation: We can pick from all 3 trees.

Example 2: Input: fruits = [0,1,2,2] Output: 3 Explanation: We can pick from trees [1,2,2]. If we had started at the first tree, we would only pick from trees [0,1].

Example 3: Input: fruits = [1,2,3,2,2] Output: 4 Explanation: We can pick from trees [2,3,2,2]. If we had started at the first tree, we would only pick from trees [1,2].

Constraints:

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

Solution

Brute Force - O(n^3) time - O(n) space

  1. Initialize max_picked = 0 to track the maximum number of fruits we can collect.
  2. Iterate over the left index left of subarrays.
  3. For every subarray start at index left, iterate over every index right to fix the end of subarray.
  4. For each subarray (left, right), count the types of fruits it contains.
    • If there are no more than 2 types, this subarray is valid, we take its length to update max_picked.
    • Otherwise, if the current subarray is invalid, we move on to the next subarray.
  5. Once we finish the iteration, return max_picked as the maximum number of fruits we can collect.
class Solution {
public:
    int totalFruit(vector<int>& fruits) {
        int maxPicked = 0;
        
        // Iterate over all subarrays (left, right)
        for (int left = 0; left < fruits.size(); ++left) {
            for (int right = 0; right < fruits.size(); ++right) {
                // Use a set to count the type of fruits.
                set<int> basket;
                
                // Iterate over the current subarray.
                for (int currentIndex = left; currentIndex <= right; ++currentIndex) {
                    basket.insert(fruits[currentIndex]);
                }
                
                // If the number of types of fruits in this subarray (types of fruits) 
                // is no larger than 2, this is a valid subarray, update 'maxPicked'.
                if (basket.size() <= 2) {
                    maxPicked = max(maxPicked, right - left + 1);
                }
            }
        }
        
        return maxPicked;
    }
};

Optimized Brute Force O(n²) time - O(1) space

  1. Initialize max_picked as 0.
  2. Iterate over left, the left index of the subarray.
  3. For every subarray start at index left, we iterate over every index right to fix the end of subarray, and calculate the types of fruits in this subarray.
    • If there are no more than 2 types, this subarray is valid, we update max_picked with the length of this subarray.
    • Otherwise, the current subarray, as well as all the longer subarrays (with the same left index left) are invalid. Move on to the next left index left + 1.
  4. Once we finish the iteration, return max_picked as the maximum number of fruits we can collect.
class Solution {
public:
    int totalFruit(vector<int>& fruits) {
        int maxPicked = 0;
        
        // Iterate over the left index left of subarrays.
        for (int left = 0; left < fruits.size(); ++left) {
            set<int> basket;
            int right = left;
            
            // Iterate over the right index right of subarrays.
            while (right < fruits.size()) {
                // Early stop. If adding this fruit makes 3 types of fruit,
                // we should stop the inner loop.
                if (basket.find(fruits[right]) == basket.end() && basket.size() == 2)
                    break;
                
                // Otherwise, update the number of this fruit.
                basket.insert(fruits[right]);
                right++;
            }
            
            // Update maxPicked.
            maxPicked = max(maxPicked, right - left);
        }
        
        // Return maxPicked as the maximum length of valid subarray.
        // (maximum number of fruits we can pick).
        return maxPicked;
    }
};

Sliding Window - O(n) time - O(n) space

  1. Start with an empty window with left and right as its left and right index.
  2. We iterate over right and add fruits[right] to this window.
    • If the number is no larger than 2, meaning that we collect no more than 2 types of fruits, this subarray is valid.
    • Otherwise, it is not the right time to expand the window and we must keep its size. Since we have added one fruit from the right side, we should remove one fruit from the left side of the window, and increment left by 1.
  3. Once we are done iterating, the difference between left and right stands for the longest valid subarray we encountered, i.e. the maximum number of fruits we can collect.
class Solution {
public:
    int totalFruit(vector<int>& fruits) {
        // Hash map 'basket' to store the types of fruits.
        unordered_map<int, int> basket;
        int left, right;
        
        // Add fruit from the right index (right) of the window.
        for (left = 0, right = 0; right < fruits.size(); ++right) {
            basket[fruits[right]]++;
            
            // If the current window has more than 2 types of fruit,
            // we remove one fruit from the left index (left) of the window.
            if (basket.size() > 2) {
                basket[fruits[left]]--;
                if (basket[fruits[left]] == 0)
                    basket.erase(fruits[left]);
                left++;
            }
        }
        
        // Once we finish the iteration, the indexes left and right 
        // stands for the longest valid subarray we encountered.
        return right - left;
    }
};

Improved Sliding Window - O(n) time - O(1) space

  1. Initialize max_picked = 0 as the maximum fruits we can collect, and use hash map basket to record the types of fruits in the current window.
  2. Start with an empty window having left = 0 and right = 0 as its left and right index.
  3. We iterate over right and add fruits[right] to this window.
    • If there are no more than 2 types of fruits, this subarray is valid.
    • Otherwise, we need to keep removing fruits from the left side until there are only 2 types of fruits in the window. Then we update max_picked as max(max_picked, right - left + 1).
  4. Once we finish iterating, return max_picked as the maximum number of fruits we can collect.
class Solution {
public:
    int totalFruit(vector<int>& fruits) {
        // We use a hash map 'basket' to store the number of each type of fruit.
        unordered_map<int, int> basket;
        int left = 0, maxPicked = 0;
        
        // Add fruit from the right index (right) of the window.
        for (int right = 0; right < fruits.size(); ++right) {
            basket[fruits[right]]++;
            
            // If the current window has more than 2 types of fruit,
            // we remove fruit from the left index (left) of the window,
            // until the window has only 2 types of fruit.
            while (basket.size() > 2) {
                basket[fruits[left]]--;
                if (basket[fruits[left]] == 0)
                    basket.erase(fruits[left]);
                left++;
            }
            
            // Update maxPicked.
            maxPicked = max(maxPicked, right - left + 1);
        }
        
        // Return maxPicked as the maximum number of fruits we can collect.
        return maxPicked;
    }
};