"""
Given an array of characters where each character represents a fruit tree, you
are given two baskets and your goal is to put maximum number of fruits in each
basket. The only restriction is that each basket can have only one type of
fruit.

Given an array of characters where each character represents a fruit tree, you
are given two baskets and your goal is to put maximum number of fruits in each
basket. The only restriction is that each basket can have only one type of fruit.

Write a function to return the maximum number of fruits in both the baskets.

Examples
========

Input: Fruit=['A', 'B', 'C', 'A', 'C']
Output: 3
Explanation: We can put 2 'C' in one basket and one 'A' in the other from the subarray ['C', 'A', 'C']

Input: Fruit=['A', 'B', 'C', 'B', 'B', 'C']
Output: 5
Explanation: We can put 3 'B' in one basket and two 'C' in the other basket.
This can be done if we start with the second letter: ['B', 'C', 'B', 'B', 'C']

Time complexity
===============
The time complexity of the above algorithm will be O(N) where ‘N’ is the number
of characters in the input array.
The outer for loop runs for all characters and the inner while loop processes
each character only once, therefore the time complexity of the algorithm will
be O(N+N) which is asymptotically equivalent to O(N)

Space Complexity
================

The algorithm runs in constant space O(1) as there can be a maximum of three
types of fruits stored in the frequency map.

Similar Problems
================

Problem 1: Longest Substring with at most 2 distinct characters

Given a string, find the length of the longest substring in it with at most two
distinct characters.

Solution: This problem is exactly similar to our parent problem.
"""

# insight: longest substring with k=2 distinct characters!


def fruits_into_basket(arr, k=2):
    window_start = max_length = 0
    char_freq = {}
    for window_end in range(len(arr)):
        right_char = arr[window_end]
        if right_char not in char_freq:
            char_freq[right_char] = 0
        char_freq[right_char] += 1

        while len(char_freq) > k:
            left_char = arr[window_start]
            char_freq[left_char] -= 1
            if char_freq[left_char] == 0:
                del char_freq[left_char]
            window_start += 1

        max_length = max(max_length, window_end - window_start + 1)
    return max_length


if __name__ == "__main__":
    arr = ["A", "B", "C", "A", "C"]
    n = fruits_into_basket(arr)
    print(f"arr: {arr}, max num fruits={n}")
    print("")
    arr = ["A", "B", "C", "B", "B", "C"]
    n = fruits_into_basket(arr)
    print(f"arr: {arr}, max num fruits={n}")
