[LeetCode] 904. Fruit Into Baskets

- Problem

904. Fruit Into Baskets

- 내 풀이 1 (Brute-force, TLE)

class Solution:
    def totalFruit(self, fruits: List[int]) -> int:
        N = len(fruits)
        answer = 0

        for i in range(N):
            baskets = dict()
            for j in range(i, N):
                if fruits[j] not in baskets:
                    baskets[fruits[j]] = 1
                else:
                    baskets[fruits[j]] += 1
                
                if len(baskets) <= 2:
                    answer = max(answer, sum(baskets.values()))

                else:
                    break
                    
        return answer

- 결과

문제에서 주어지는 인자 fruits의 범위는 1~10**5 이다.
위의 코드는 시간 복잡도 O(N^2) 이므로 당연히 TLE가 날 수밖에 없다.

- 내 풀이 2 (hash table)

class Solution:
    def totalFruit(self, fruits: List[int]) -> int:
        baskets = dict()
        answer = prev = 0

        for fruit in fruits:
            if fruit in baskets:
                baskets[fruit] += 1
            
            else:
                baskets[fruit] = 1
            
            if len(baskets) > 2:
                baskets[fruits[prev]] -= 1
                if not baskets[fruits[prev]]:
                    baskets.pop(fruits[prev])
                prev += 1

            else:
                answer = max(answer, sum(baskets.values()))

        return answer

- 결과

• 시간 복잡도: O(N) (N은 주어진 fruits의 길이)
• 공간 복잡도: O(1)
  - baksets 은 최대 2가지의 키:밸류 값만 가질 수 있다.