Fruit Into Baskets (LeetCode 904)

The problem

A row of trees is given as an array fruits, where fruits[i] is the type of fruit on tree i. You have two baskets, and each basket holds only one type of fruit (as much as you like). Starting at any tree, you pick one fruit from every tree moving right, and you must stop when you reach a tree whose type fits in neither basket. Return the most fruit you can pick.

Input:  fruits = [1, 2, 3, 2, 2]
Output: 4

Start at index 1 and pick [2, 3, 2, 2]: only types 2 and 3, four fruits.

How to think about it

Strip away the story. You pick a contiguous stretch of trees, and it may contain at most two distinct values. So the problem is really: find the longest subarray with at most 2 distinct numbers.

Start with brute force. From each start index, walk right while tracking the distinct types in a set, and stop on the third type. That’s O(n²) in the worst case.

The key insight: the condition is monotonic, so use a variable window. If a window has more than two types, any window that contains it also has more than two. So once the window breaks the rule, the only fix is to move the left edge right until one type disappears. Keep a count map of the types in the window; a type leaves the window when its count drops to zero, and the map’s size tells you how many distinct types are inside.

Which pattern is this? A variable-size sliding window with a frequency map, often called “longest substring with at most K distinct characters” (here K = 2). The tell is “longest contiguous stretch with at most K kinds”.

Approach

1. Keep cnt (type → count in the window), start = 0, best = 0.
2. For each i, add fruits[i] to cnt.
3. While cnt has more than two keys, decrement fruits[start], delete it if its count hits 0, and move start right.
4. Update best = max(best, i - start + 1).
5. Return best.

Solution

from collections import defaultdict

def total_fruit(fruits):
    cnt = defaultdict(int); start = 0; best = 0
    for i, f in enumerate(fruits):
        cnt[f] += 1
        while len(cnt) > 2:                               # three types: shrink
            cnt[fruits[start]] -= 1
            if cnt[fruits[start]] == 0: del cnt[fruits[start]]
            start += 1
        best = max(best, i - start + 1)
    return best
# total_fruit([1,2,3,2,2]) -> 4

Dry run

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

At i = 2, type 3 makes three kinds, so type 1 is dropped. The answer is 4.

Complexity

• Time: O(n). Each index is added once and removed at most once.
• Space: O(1). The map never holds more than three keys.

Edge cases and pitfalls

• One tree returns 1. Only one or two types overall returns len(fruits).
• Forgetting to delete zero-count keys means len(cnt) never shrinks, and the while loop runs off the end.
• Resetting the window to the last tree when a third type arrives is wrong. In [1, 2, 1, 3], the best window after the 3 still includes the 1 before it.
• Using a set instead of counts can’t tell you when a type has fully left the window.
• Generalising to K types is just changing 2 to K.

Related problems

• Longest Substring Without Repeating Characters (LeetCode 3): the same longest-window template with a “no repeats” rule.
• Max Consecutive Ones III (LeetCode 1004): a window limited by a count (zeros) rather than distinct values.
• Longest Repeating Character Replacement (LeetCode 424): a window validated by a frequency budget.