LeetCode 904: Fruit Into Baskets

Problem Restatement

We are given an integer array fruits.

Each value represents the fruit type of one tree in a row.

We have two baskets. Each basket can hold only one fruit type, but it can hold unlimited fruits of that type.

Starting from any tree, we must move right and pick exactly one fruit from every tree. We cannot skip trees.

We stop when we reach a tree whose fruit type does not fit in either basket.

Return the maximum number of fruits we can pick.

This is the same as asking:

Find the length of the longest contiguous subarray with at most two distinct values.

Input and Output

Function shape:

def totalFruit(fruits: list[int]) -> int:
    ...

Examples

Example 1:

fruits = [1, 2, 1]

We can pick all trees because there are only two fruit types.

Answer:

3

Example 2:

fruits = [0, 1, 2, 2]

The best valid segment is:

[1, 2, 2]

It has two fruit types and length 3.

Answer:

3

Example 3:

fruits = [1, 2, 3, 2, 2]

The best valid segment is:

[2, 3, 2, 2]

It has fruit types 2 and 3.

Answer:

4

First Thought: Check Every Starting Tree

A direct method is to try every possible starting index.

From each start, move right while the segment contains at most two fruit types.

class Solution:
    def totalFruit(self, fruits: list[int]) -> int:
        n = len(fruits)
        best = 0

        for start in range(n):
            types = set()

            for end in range(start, n):
                types.add(fruits[end])

                if len(types) > 2:
                    break

                best = max(best, end - start + 1)

        return best

This works, but it repeats too much work.

For each starting index, we scan forward again. In the worst case, this can take O(n^2) time.

Key Insight

We need the longest contiguous window that contains at most two distinct fruit types.

This is a standard sliding window condition.

We keep a window:

fruits[left:right + 1]

The window is valid when it contains at most two distinct fruit types.

As we move right from left to right, we add one fruit to the window.

If the window has more than two fruit types, we move left rightward until the window becomes valid again.

During the scan, every valid window is a possible answer.

Algorithm

Use a hash map called count.

It stores how many times each fruit type appears in the current window.

Steps:

1. Set left = 0.
2. Set answer = 0.
3. Iterate right from 0 to len(fruits) - 1.
4. Add fruits[right] to count.
5. While count has more than two fruit types:
   • Remove fruits[left] from the window.
   • If its count becomes 0, delete that fruit type.
   • Move left one step right.
6. Update the answer with the current window length.

Correctness

The algorithm keeps the window valid after each iteration.

When a new fruit is added, the window may contain three fruit types. If so, moving left rightward removes fruits from the left side until only two types remain.

After this shrinking step, the window ending at right is the longest valid window ending at right. If we moved left any less, the window would still have more than two fruit types. If we moved it more, the window would become shorter without need.

The algorithm checks the best valid window ending at every possible right index.

Since every valid contiguous picking sequence has some ending index, it will be considered when that ending index is processed.

Therefore, the maximum length recorded by the algorithm is exactly the maximum number of fruits we can pick.

Complexity

Although the hash map is used, the basket limit is fixed at two types, so the space is constant.

Implementation

from collections import defaultdict

class Solution:
    def totalFruit(self, fruits: list[int]) -> int:
        count = defaultdict(int)
        left = 0
        answer = 0

        for right, fruit in enumerate(fruits):
            count[fruit] += 1

            while len(count) > 2:
                left_fruit = fruits[left]
                count[left_fruit] -= 1

                if count[left_fruit] == 0:
                    del count[left_fruit]

                left += 1

            answer = max(answer, right - left + 1)

        return answer

Code Explanation

The hash map stores fruit counts inside the current window:

count = defaultdict(int)

The left boundary starts at the first tree:

left = 0

We expand the window by moving right:

for right, fruit in enumerate(fruits):
    count[fruit] += 1

If there are more than two fruit types, the window breaks the basket rule:

while len(count) > 2:

So we remove fruits from the left:

left_fruit = fruits[left]
count[left_fruit] -= 1

If a fruit type no longer appears in the window, remove it from the map:

if count[left_fruit] == 0:
    del count[left_fruit]

Then move the left boundary:

left += 1

After the window is valid again, update the best answer:

answer = max(answer, right - left + 1)

Testing

def run_tests():
    s = Solution()

    assert s.totalFruit([1, 2, 1]) == 3
    assert s.totalFruit([0, 1, 2, 2]) == 3
    assert s.totalFruit([1, 2, 3, 2, 2]) == 4
    assert s.totalFruit([3, 3, 3, 1, 2, 1, 1, 2, 3, 3, 4]) == 5
    assert s.totalFruit([1]) == 1
    assert s.totalFruit([1, 1, 1, 1]) == 4
    assert s.totalFruit([1, 2, 3, 4, 5]) == 2

    print("all tests passed")

run_tests()

Test meaning: