Leetcode 904 - Fruit Into Baskets

Understanding the problem

You are given an array tree of integers, where tree[i] is the type of a fruit tree planted at the ith position.

You are also given two integers max_tree and max_fruit, where max_tree is the maximum number of trees that can be planted in a single basket, and max_fruit is the maximum number of fruits that can be picked from a single basket.

Return the maximum number of fruits that you can pick from the tree.

Example:

Input: tree = [1, 2, 1]

Output: 3

Explanation: We can pick at most 3 fruits from the tree. We can pick the fruits from the first basket (1, 2) or from the second basket (2, 1).

Input: tree = [0, 1, 2, 2]

Output: 4

Explanation: We can pick at most 4 fruits from the tree. We can pick the fruits from the first two baskets (0, 1) or (1, 2).

Constraints:

• 1 <= tree.length <= 10^4
• 0 <= tree[i] <= 2
• 1 <= max_tree <= 20
• 1 <= max_fruit <= 10^5

Plan your solution

One approach to solve this problem is to use a sliding window.

We can start by initializing two variables start and end to 0. These variables will keep track of the start and end indices of the current window, respectively.

We can also initialize two variables tree_count and fruit_count to 0. These variables will keep track of the number of different trees and the number of fruits in the current window, respectively.

We can also initialize a variable max_fruit_count to 0. This variable will keep track of the maximum number of fruits that can be picked from the tree.

Then, we can use a while loop to iterate over the elements of fruits and do the following:

• If the current element is not in the current window, increment tree_count by 1.
• Increment fruit_count by 1.
• If tree_count is greater than 2, remove the first element of the current window and decrement the counts accordingly.
• Update max_fruit_count with the maximum of max_fruit_count and fruit_count.
• Increment end by 1.

Finally, we return max_fruit_count.

Implement your solution

Here is the complete implementation:

class Solution:
    def totalFruit(self, fruits: List[int]) -> int:
        # Initialize start, end, tree_count, fruit_count and max_fruit_count
        start = end = tree_count = fruit_count = 0
        max_fruit_count = 0
        
        # Iterate over elements of fruits
        while end < len(fruits):
            # If the current element is not in the current window, increment tree_count by 1
            if fruits[end] not in fruits[start:end]:
                tree_count += 1
            # Increment fruit_count by 1
            fruit_count += 1
            # If tree_count is greater than 2, remove the first element of the current window and decrement the counts accordingly
            while tree_count > 2:
                start += 1
                if fruits[start-1] in fruits[start:end]:
                    fruit_count -= 1
                else:
                    tree_count -= 1
            # Update max_fruit_count with the maximum of max_fruit_count and fruit_count
            max_fruit_count = max(max_fruit_count, fruit_count)
            # Increment end by 1
            end += 1
        
        # Return max_fruit_count
        return max_fruit_count

Test your solution

Now you can test your solution with some test cases.

fruits = [1, 2, 1]
assert totalFruit(fruits) == 3

fruits = [0, 1, 2, 2]
assert totalFruit(fruits) == 4

fruits = [1, 2, 3, 2, 2]
assert totalFruit(fruits) == 4

Related:

• permutation in string
• two sum solution with code
• kth smallest element in a bst with code
• replace words with solution