904. Fruit Into Baskets

Reformulated question¶

Find the length of the longest contiguous subarray containing at most 2 distinct values.

Compact example:

fruits = [1,2,3,2,2] -> 4
best window: [2,3,2,2]

Key trick¶

Use a sliding window.

Expand the right end.
Count fruit types in the current window.
While there are more than 2 types, move the left end rightward.
Track the maximum window length.

Trap¶

Treating it like a subsequence instead of a contiguous subarray.
Restarting from many positions, which becomes \(O(n^2)\).
Forgetting to remove a fruit type from the counter when its count becomes 0.
Missing that "start anywhere, move only right" means "choose any contiguous segment".

Why this question is interesting¶

It is a clean "at most K distinct" sliding-window pattern.

Very common interview template.
Simple to code.
Tests correctness and window-shrinking intuition.

Solve the problem with idiomatic python¶

class Solution:
    def totalFruit(self, fruits: list[int]) -> int:
        # Sliding window with counts of fruit types inside the window.
        counts = {}
        left = 0
        best = 0

        for right, fruit in enumerate(fruits):
            counts[fruit] = counts.get(fruit, 0) + 1

            # Keep only at most 2 distinct fruit types.
            while len(counts) > 2:
                left_fruit = fruits[left]
                counts[left_fruit] -= 1
                if counts[left_fruit] == 0:
                    del counts[left_fruit]
                left += 1

            # Window [left, right] is valid.
            best = max(best, right - left + 1)

        return best

Pytest test¶

import pytest

from solution import Solution


@pytest.mark.parametrize(
    ("fruits", "expected"),
    [
        ([1], 1),
        ([1, 2, 1], 3),
        ([0, 1, 2, 2], 3),
        ([1, 2, 3, 2, 2], 4),
        ([3, 3, 3, 1, 2, 1, 1, 2, 3, 3, 4], 5),
        ([1, 1, 1, 1], 4),
        ([1, 2, 1, 2, 1, 2], 6),
        ([1, 2, 3, 4], 2),
        ([0, 1, 6, 6, 4, 4, 6], 5),
    ],
)
def test_total_fruit(fruits, expected):
    assert Solution().totalFruit(fruits) == expected

Comment my solution¶

Your solution is correct on many cases, but it is not a good interview answer.

It explores multiple starting points with a stack, so worst-case time is \(O(n^2)\).
The intended solution is \(O(n)\) with one sliding window.
Using -1 as a sentinel for t2 is brittle even if constraints make it safe.
The stack.append((i, t2)) branching is clever, but harder to reason about than the standard window pattern.

A counterexample for performance:

fruits = [1,2,3,4,5,6,...]

Your code keeps retrying from many indices.
Sliding window processes each index at most twice.

## Solution
# WRONG
class Solution:
    def totalFruit(self, fruits: list[int]) -> int:
        stack = [(0, fruits[0])]
        best_count = 0
        while stack:
            i, t1 = stack.pop()
            count = 0
            t2 = -1
            while i < len(fruits):
                if fruits[i] == t1:
                    count += 1
                elif fruits[i] == t2:
                    count += 1
                elif t2 == -1:
                    t2 = fruits[i]
                    stack.append((i, t2))
                    count += 1
                else:
                    break
                i += 1
            best_count = max(best_count, count)
        return best_count

# WRONG
# written after reading AI comments
class Solution:
    def totalFruit(self, fruits: list[int]) -> int:
        max_window_length = 0
        left, right = 0, 0
        while right < len(fruits):
            right += 1
            seen = set()
            for i in range(left, right):
                seen.add(fruits[i])
            if len(seen) <= 2:
                max_window_length = max(max_window_length, right - left)
                continue
            while len(seen) > 2:
                left += left
                seen = set()
                for i in range(left, right):
                    seen.add(fruits[i])
        return max_window_length