LeetCode 904 - Fruit Into Baskets

This problem asks us to find the length of the longest contiguous subarray that contains at most two distinct values. Each element in the fruits array represents the type of fruit produced by a tree.

leetcodemediumarrayhash-tablesliding-window

LeetCode Problem 904

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Sliding Window

Solution

Problem Understanding

This problem asks us to find the length of the longest contiguous subarray that contains at most two distinct values.

Each element in the fruits array represents the type of fruit produced by a tree. We are allowed to start at any position and move only to the right, collecting exactly one fruit from every tree we visit. However, we only have two baskets, and each basket can hold only one fruit type.

That means that while traversing the array, we may collect fruits from at most two distinct fruit types. The moment we encounter a third fruit type, we must stop.

The goal is to determine the maximum number of fruits we can collect in one continuous walk.

For example, in:

[1,2,3,2,2]

the longest valid contiguous segment with at most two distinct fruit types is:

[2,3,2,2]

which has length 4.

The constraints are important:

  • 1 <= fruits.length <= 10^5
  • 0 <= fruits[i] < fruits.length

An input size of 10^5 means an O(n^2) solution will likely be too slow. We need a linear or near-linear algorithm.

The problem guarantees:

  • The array is non-empty.
  • Fruit types are integers.
  • We only care about contiguous segments because we must continuously move to the right without skipping trees.

Several edge cases are important to consider:

  • Arrays containing only one fruit type.
  • Arrays with exactly two fruit types.
  • Arrays where every element is different.
  • Situations where the valid window must shrink multiple times after encountering a third fruit type.
  • Very small arrays such as length 1.

A naive implementation may repeatedly scan subarrays, leading to quadratic time complexity.

Approaches

Brute Force Approach

The brute force idea is to consider every possible starting position and extend the subarray to the right until more than two distinct fruit types appear.

For each starting index:

  1. Create a set or frequency map.
  2. Extend the ending index one step at a time.
  3. Stop when the number of distinct fruit types exceeds two.
  4. Record the maximum valid length found.

This works because it explicitly checks every possible contiguous segment.

However, the problem is efficiency. If we try all starting positions and scan forward each time, the total complexity becomes O(n^2) in the worst case.

With n = 10^5, quadratic complexity is too slow.

Optimal Sliding Window Approach

The key observation is that we only care about contiguous subarrays containing at most two distinct fruit types.

This is a classic sliding window scenario:

  • We maintain a window [left, right].
  • The window always satisfies the rule:
at most two distinct fruit types

As we expand the window by moving right forward:

  • We add the new fruit into a frequency map.
  • If the window becomes invalid, meaning it contains more than two fruit types, we shrink it from the left until it becomes valid again.

At every step, the current window represents a valid collection path.

Because each element enters and leaves the window at most once, the algorithm runs in linear time.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(1) to O(2) Tries every starting point and expands forward
Optimal Sliding Window O(n) O(1) to O(3) Maintains a valid window with at most two fruit types

Algorithm Walkthrough

Optimal Sliding Window Algorithm

  1. Initialize two pointers:
  • left = 0
  • right will iterate through the array.
  1. Create a hash map called fruit_count.

This map stores:

fruit type -> frequency inside current window
  1. Initialize a variable max_fruits = 0.
  2. Iterate through the array using right.
  3. For each fruits[right]:
  • Add it to the frequency map.
  • Increase its count.
  1. After adding the new fruit, check whether the window has become invalid.

The window is invalid when:

len(fruit_count) > 2
  1. While the window is invalid:
  • Remove fruits[left] from the window.
  • Decrease its count.
  • If its count becomes zero, remove that fruit type entirely from the map.
  • Move left forward.
  1. At this point, the window is valid again.

The current window length is:

right - left + 1
  1. Update max_fruits if the current window is larger.
  2. Continue until the entire array has been processed.
  3. Return max_fruits.

Why it works

The sliding window always maintains a valid contiguous subarray containing at most two distinct fruit types.

Whenever a third fruit type appears, we shrink the window from the left until only two types remain again.

Because every valid window is considered exactly once during expansion, and because the algorithm always keeps the largest valid window seen so far, the final answer is guaranteed to be correct.

Python Solution

from typing import List
from collections import defaultdict

class Solution:
    def totalFruit(self, fruits: List[int]) -> int:
        left = 0
        max_fruits = 0
        
        fruit_count = defaultdict(int)
        
        for right in range(len(fruits)):
            fruit = fruits[right]
            fruit_count[fruit] += 1
            
            while len(fruit_count) > 2:
                left_fruit = fruits[left]
                fruit_count[left_fruit] -= 1
                
                if fruit_count[left_fruit] == 0:
                    del fruit_count[left_fruit]
                
                left += 1
            
            current_window = right - left + 1
            max_fruits = max(max_fruits, current_window)
        
        return max_fruits

The implementation begins by initializing the sliding window boundaries and a frequency map.

The fruit_count dictionary tracks how many times each fruit type appears inside the current window. This allows us to efficiently determine how many distinct fruit types are currently present.

The right pointer expands the window one element at a time. Every new fruit is added into the frequency map.

If adding the new fruit causes the number of distinct fruit types to exceed two, the window becomes invalid. The while loop then shrinks the window from the left side until the constraint is satisfied again.

Whenever a fruit count becomes zero, that fruit type is removed from the dictionary entirely. This ensures that:

len(fruit_count)

correctly represents the number of distinct fruit types in the window.

After restoring validity, the current window length is computed and compared against the best answer found so far.

Because each pointer only moves forward, the algorithm runs efficiently in linear time.

Go Solution

func totalFruit(fruits []int) int {
    left := 0
    maxFruits := 0

    fruitCount := make(map[int]int)

    for right := 0; right < len(fruits); right++ {
        fruit := fruits[right]
        fruitCount[fruit]++

        for len(fruitCount) > 2 {
            leftFruit := fruits[left]
            fruitCount[leftFruit]--

            if fruitCount[leftFruit] == 0 {
                delete(fruitCount, leftFruit)
            }

            left++
        }

        currentWindow := right - left + 1

        if currentWindow > maxFruits {
            maxFruits = currentWindow
        }
    }

    return maxFruits
}

The Go implementation follows the same sliding window logic as the Python solution.

Go uses a built-in map:

map[int]int

to store fruit frequencies.

Unlike Python's defaultdict, Go maps return the zero value automatically for missing keys, so incrementing counts works naturally.

The delete() function removes fruit types whose counts become zero, ensuring the map size accurately reflects the number of distinct fruit types in the current window.

Integer overflow is not a concern because the maximum array length is only 10^5.

Worked Examples

Example 1

Input:

[1,2,1]
Step right Fruit Window fruit_count max_fruits
1 0 1 [1] {1:1} 1
2 1 2 [1,2] {1:1,2:1} 2
3 2 1 [1,2,1] {1:2,2:1} 3

Final answer:

3

Example 2

Input:

[0,1,2,2]
Step right Fruit Window fruit_count max_fruits
1 0 0 [0] {0:1} 1
2 1 1 [0,1] {0:1,1:1} 2
3 2 2 Invalid {0:1,1:1,2:1} 2

Now shrink:

left movement Window fruit_count
Remove 0 [1,2] {1:1,2:1}

Continue:

Step right Fruit Window fruit_count max_fruits
4 3 2 [1,2,2] {1:1,2:2} 3

Final answer:

3

Example 3

Input:

[1,2,3,2,2]
Step right Fruit Window fruit_count max_fruits
1 0 1 [1] {1:1} 1
2 1 2 [1,2] {1:1,2:1} 2
3 2 3 Invalid {1:1,2:1,3:1} 2

Shrink window:

left movement Window fruit_count
Remove 1 [2,3] {2:1,3:1}

Continue:

Step right Fruit Window fruit_count max_fruits
4 3 2 [2,3,2] {2:2,3:1} 3
5 4 2 [2,3,2,2] {2:3,3:1} 4

Final answer:

4

Complexity Analysis

Measure Complexity Explanation
Time O(n) Each element enters and leaves the window at most once
Space O(1) The map stores at most three fruit types

The algorithm is linear because both pointers move only forward. Even though there is a nested while loop, the total number of left pointer movements across the entire algorithm is at most n.

The space complexity is effectively constant because the frequency map never stores more than three distinct fruit types at any time.

Test Cases

sol = Solution()

assert sol.totalFruit([1,2,1]) == 3  # basic example with two fruit types
assert sol.totalFruit([0,1,2,2]) == 3  # requires shrinking window
assert sol.totalFruit([1,2,3,2,2]) == 4  # optimal window appears later
assert sol.totalFruit([1]) == 1  # single element array
assert sol.totalFruit([1,1,1,1]) == 4  # all same fruit type
assert sol.totalFruit([1,2]) == 2  # exactly two fruit types
assert sol.totalFruit([1,2,3]) == 2  # third type immediately breaks window
assert sol.totalFruit([3,3,3,1,2,1,1,2,3,3,4]) == 5  # classic sliding window stress case
assert sol.totalFruit([0,1,6,6,4,4,6]) == 5  # multiple shrink operations
assert sol.totalFruit([1,0,1,4,1,4,1,2,3]) == 5  # alternating valid window
assert sol.totalFruit([1,2,1,2,1,2,1]) == 7  # entire array valid
assert sol.totalFruit([1,2,3,4,5]) == 2  # all distinct values

Test Summary

Test Why
[1,2,1] Small valid example
[0,1,2,2] Requires shrinking after third fruit
[1,2,3,2,2] Longest window occurs in the middle
[1] Minimum input size
[1,1,1,1] Only one fruit type
[1,2] Exactly two types
[1,2,3] Third type immediately invalidates window
[3,3,3,1,2,1,1,2,3,3,4] Stress test for sliding window updates
[0,1,6,6,4,4,6] Repeated shrinking and expansion
[1,0,1,4,1,4,1,2,3] Alternating valid patterns
[1,2,1,2,1,2,1] Entire array remains valid
[1,2,3,4,5] Every element distinct

Edge Cases

Single Element Array

An array containing only one tree is the smallest valid input.

Example:

[5]

A buggy implementation might incorrectly assume at least two elements exist or mishandle window initialization.

This implementation handles it naturally because the sliding window starts empty and expands normally. The maximum window size becomes 1.

All Trees Have the Same Fruit Type

Example:

[2,2,2,2,2]

Some implementations unnecessarily shrink the window or miscount distinct fruit types.

Here, the frequency map always contains exactly one key, so the window never becomes invalid. The algorithm correctly returns the entire array length.

Every Fruit Type Is Different

Example:

[1,2,3,4,5]

This case forces the window to shrink frequently.

A naive implementation may fail to properly remove fruit types whose counts become zero, causing the distinct count to remain incorrect.

This solution deletes entries from the frequency map whenever their frequency reaches zero, ensuring the window validity condition remains accurate.

Long Alternating Pattern

Example:

[1,2,1,2,1,2,1]

The entire array is valid because only two fruit types appear.

An incorrect implementation might prematurely shrink the window when seeing repeated alternation.

This algorithm correctly recognizes that the number of distinct fruit types never exceeds two, so the entire array is counted.