904.fruit-into-baskets.md

904. Fruit Into Baskets

Sliding Window

The problem is asking: "Find the longest subarray that contains at most two types of fruits". We maintain a window that contains at most two types of fruits, we try to expand the window if the number of type <= 2, otherwise we shrink the window from the left side.

Please note that we have to remove the key from the map if the count is 0. If we don't remove the key when the count drops to 0, then the map key size is still 3, even if only 2 types of fruits are in the window.

def totalFruit(self, fruits: List[int]) -> int:
    left, right = 0, 0
    baskets = {}
    pick = 0
    while right < len(fruits):
        baskets[fruits[right]] = baskets.get(fruits[right], 0) + 1
        while len(baskets.keys()) > 2:
            baskets[fruits[left]] -= 1
            if baskets[fruits[left]] == 0: del baskets[fruits[left]]
            left += 1
        pick = max(pick, right - left + 1)
        right += 1
    return pick

fun totalFruit(fruits: IntArray): Int {
    var ans = 0
    var left = 0
    val map = HashMap<Int, Int>()
    for (right in fruits.indices) {
        map[fruits[right]] = (map[fruits[right]] ?: 0) + 1
        while (map.size > 2) {
            val count = map[fruits[left]]!! - 1
            if (count > 0) map[fruits[left]] = count
            else map.remove(fruits[left])
            left++
        }
        ans = maxOf(ans, right - left + 1)
    }
    return ans
}

In this problem, we can't use set to track the types of fruits:

val set = mutableSetOf<Int>()
var left = 0
for (right in fruits.indices) {
    set.add(fruits[right])

    while (set.size > 2) {
        set.remove(fruits[left])  // ❌ invalid removal
        left++
    }
}

For example, [1, 2, 3, 1], and we will remove 1 from the set, but we still have 1 (the first element) in the set.

• Time Complexity: O(n).
• Space Complexity: O(1), there are at most 3 types of fruits in the baskets.