Fruit Into Baskets Python Solution with Tests

lcpy gen -n 904   # by problem number
lcpy gen -s fruit_into_basket   # by problem name

Input: fruits = [1,2,1]
Output: 3
Explanation: We can pick from all 3 trees.

Input: fruits = [0,1,2,2]
Output: 3
Explanation: We can pick from trees [1,2,2].
If we had started from the first tree, we would only pick from trees [0,1].

Input: fruits = [1,2,3,2,2]
Output: 4
Explanation: We can pick from trees [2,3,2,2].
If we had started from the first tree, we would only pick from trees [1,2].

from collections import defaultdict


class Solution:
    # Time: O(n)
    # Space: O(1)
    def total_fruit(self, fruits: list[int]) -> int:
        counts: defaultdict[int, int] = defaultdict(int)
        left = 0
        best = 0
        for right, fruit in enumerate(fruits):
            counts[fruit] += 1
            while len(counts) > 2:
                left_fruit = fruits[left]
                counts[left_fruit] -= 1
                if counts[left_fruit] == 0:
                    del counts[left_fruit]
                left += 1
            best = max(best, right - left + 1)
        return best