Fruit Into Baskets LeetCode solution in Python

1
class Solution:
2
    def totalFruit(self, fruits: List[int]) -> int:
3
        ans = 0
4
        fruitdict = defaultdict()
5
        stack = []
6
        i, j = 0, 0
7

8
        while j < len(fruits):
9
            if fruits[j] not in fruitdict and len(fruitdict) < 2:
10
                stack.append(fruits[j])
11
                fruitdict[fruits[j]] = j
12
                j += 1
13

14
            elif fruits[j] in fruitdict:
15
                fruitdict[fruits[j]] = j
16
                j += 1
17

18
            else:
19
                if fruitdict[stack[0]] > fruitdict[stack[1]]:
20
                    i = fruitdict[stack[1]] + 1
21
                    del fruitdict[stack[1]]
22
                    stack.pop()
23
                else:
24
                    i = fruitdict[stack[0]] + 1
25
                    del fruitdict[stack[0]]
26
                    stack.pop(0)
27

28
            ans = max(ans, j - i)
29
        return ans