How to solve the Fruit Into Baskets problem in Python

Problem statement

Given a list of fruit types (in the form of integers), where each integer corresponds to a specific tree, our goal is to find the maximum number of fruits we can pick while following these constraints:

We have two baskets. Each basket can hold only one type of fruit, but there is no limit to the number of fruits it can hold.
We pick one fruit from each tree while moving toward the right. The fruits that we pick have to fit in our baskets.
We must stop once we reach a fruit that doesn’t fit in the baskets.

Let’s take a look at an example below:

Algorithm

In this code, we’ll maintain a window of contiguous trees to contain at most two fruits. We’ll use the concept of sliding windows. We’ll also use a dictionary to keep track of the count of each fruit type and a few pointers to manage the boundaries of our window, which are as follows:

left: It indicates the start of the current window. It moves to the right whenever the condition of having more than two types of fruits in the window is met.
right: It indicates the end of the current window.

The maximum number of fruits we can pick is the maximum window size achieved while following the constraints.

def totalFruit(tree):
    maximum = 0      # Longest subarray length seen so far
    left = 0         # Left index of the current subarray (window)
    fruitscounter = {}  # Counts occurrences of fruits in the current window
    for right, fruit in enumerate(tree):
        fruitscounter[fruit] = fruitscounter.get(fruit, 0) + 1  # Add/update fruit count
        # If we have more than two types of fruits in the window, shrink it from the left
        while len(fruitscounter) > 2:
            fruitscounter[tree[left]] -= 1 
            if fruitscounter[tree[left]] == 0:
                del fruitscounter[tree[left]]
            left += 1
        # Update the maximum length if the current window is longer
        maximum = max(maximum, right - left + 1) 
    return maximum
# Examples
print(totalFruit([1, 2, 1, 2, 3]))  # Output: 4
print(totalFruit([3, 3, 3, 1, 2, 1, 1, 2, 3, 3, 4]))  # Output: 5

Explanation

Line 1: This is the definition of the totalFruit function, which takes tree as a parameter.
Lines 2–4: We set the maximum and left values to 0, and make an empty dictionary called fruitscounter that holds the count of each fruit.
Line 6: We iterate through the tree list using an enumeration and a for loop to process each fruit in the list. We use this to loop through the several windows.
Line 7: We update the fruitscounter to count the frequency of the current fruit type in the current window.
Line 10: We use a while loop that runs until it encounters two types of fruits in the window.
Line 11: We get the type of fruit at the left index and decrease the count in the fruitscounter dictionary.
Line 12: We check if the count of fruit at the left index has reached 0. If it does, this means that there are no more fruits of the type of index left.
Line 13: We delete the entry of the fruit type from the fruitscounter dictionary.
Line 14: We move the left pointer one position to the right.
Line 17: We store the maximum amount of fruits in the current window that can be picked with the use of maximum.
Lines 23–24: We code to run our function.

Time complexity

The time complexity of this algorithm is $O(n)$ , where n is the number of fruits in the input list, tree. The code uses a single loop to iterate through all the fruits exactly once.

Space complexity

The space complexity of the code is also $O(n)$ . This is primarily due to the size of the dictionary fruitscounter , which, in the worst-case scenario, could store a frequency count for each unique fruit in the input list if all fruits differ. However, in practical terms, since the problem is constrained to finding the longest subarray with at most two distinct integers, the space complexity can effectively be seen as $O(2)$ or $O(1)$ , as there will never be more than two unique fruits in the counter at any time.

Free Resources

Learn in-demand tech skills in half the time

PRODUCTS

Mock Interview

New

Courses

Skill Paths

Projects

Assessments