fruit-into-baskets.py

class Solution:
    def totalFruit(self, tree):
        """
        :type tree: List[int]
        :rtype: int
        """
        from collections import Counter
        count = Counter()
        
        i = 0
        ans = 0
        for j, x in enumerate(tree):
            count[x] += 1
            while len(count) > 2:
                count[tree[i]] -= 1
                if count[tree[i]] == 0:
                    del count[tree[i]]
                i += 1
            ans = max(ans, j - i + 1)
        return ans