Fruit Into Baskets [Medium] — Longest Subarray with 2 Distinct

Problem Statement

You have fruits[i] on each tree and two baskets (each holds one type). Starting from any tree, pick consecutive fruits, at most one type per basket. Return max fruits you can collect.

This is equivalent to: longest subarray with at most 2 distinct values.

Input: [1,2,1] → Output: 3
Input: [0,1,2,2] → Output: 3 ([1,2,2])
Input: [1,2,3,2,2] → Output: 4 ([2,3,2,2])

Intuition

Sliding window with a HashMap tracking counts of at most 2 fruit types. When 3rd type enters, shrink from left.

Solutions

C++

int totalFruit(vector<int>& fruits) {
    unordered_map<int,int> cnt;
    int left=0, ans=0;
    for(int right=0;right<fruits.size();right++){
        cnt[fruits[right]]++;
        while(cnt.size()>2){
            if(--cnt[fruits[left]]==0) cnt.erase(fruits[left]);
            left++;
        }
        ans=max(ans,right-left+1);
    }
    return ans;
}

Java

public int totalFruit(int[] fruits) {
    Map<Integer,Integer> cnt=new HashMap<>();
    int left=0, ans=0;
    for(int right=0;right<fruits.length;right++){
        cnt.merge(fruits[right],1,Integer::sum);
        while(cnt.size()>2){
            cnt.merge(fruits[left],-1,Integer::sum);
            if(cnt.get(fruits[left])==0) cnt.remove(fruits[left]);
            left++;
        }
        ans=Math.max(ans,right-left+1);
    }
    return ans;
}

Python

from collections import defaultdict

def totalFruit(fruits: list[int]) -> int:
    cnt = defaultdict(int)
    left = ans = 0
    for right, f in enumerate(fruits):
        cnt[f] += 1
        while len(cnt) > 2:
            cnt[fruits[left]] -= 1
            if cnt[fruits[left]] == 0:
                del cnt[fruits[left]]
            left += 1
        ans = max(ans, right - left + 1)
    return ans

Complexity

Time: O(n)
Space: O(1) — at most 3 entries in map

Generalization

Replace > 2 with > k for "at most k distinct values" variant.