# In a row of trees, the i-th tree produces fruit with type tree[i].
#
# You start at any tree of your choice, then repeatedly perform the following steps:
#
# Add one piece of fruit from this tree to your baskets.  If you cannot, stop.
# Move to the next tree to the right of the current tree.  If there is no tree to the right, stop.
# Note that you do not have any choice after the initial choice of starting tree: you must perform step 1,
# then step 2, then back to step 1, then step 2, and so on until you stop.
#
# You have two baskets, and each basket can carry any quantity of fruit, but you want each basket to only carry one
# type of fruit each.
#
# What is the total amount of fruit you can collect with this procedure?
#
#
#
# Example 1:
#
# Input: [1,2,1]
# Output: 3
# Explanation: We can collect [1,2,1].
# Example 2:
#
# Input: [0,1,2,2]
# Output: 3
# Explanation: We can collect [1,2,2].
# If we started at the first tree, we would only collect [0, 1].


# this question is exactly the same as 159 longest substring with at most 2 distinct characters

class Solution(object):
    def totalFruit(self, tree):
        """
        :type tree: List[int]
        :rtype: int
        """
        res, start, distince_count, visited = 0, 0, 0, [0] * len(tree)
        for i, c in enumerate(tree):
            if visited[c] == 0:
                distince_count += 1
            visited[c] += 1

            while distince_count > 2:
                visited[tree[start]] -= 1
                if visited[tree[start]] == 0:
                    distince_count -= 1

                start += 1

            res = max(res, i - start + 1)

        return res


if __name__ == '__main__':
    print Solution().totalFruit([1, 2, 1])
    print Solution().totalFruit([0, 1, 2, 2])
    print Solution().totalFruit([1, 2, 3, 2, 2])
    print Solution().totalFruit([3, 3, 3, 1, 2, 1, 1, 2, 3, 3, 4])
