Leetcode 904 Fruit into Buskets（Python 实现）

题目描述

题目相当于在一个数组 tree中找出一个最长的子数组，要求这个子数组中只有两种数。（数组名为 tree与题干中的背景有关。）

算法

对 tree进行扫描，记录扫描到第 i个数时，以当前数为结尾的最大合法子数组的长度 cur_len，以及元素全都与当前数相同的最大子数组的长度 same_len，以及对应的 2 个数（也可能是 1 个，特殊处理）用集合存储 type_set。当扫描下一个数时，如果发现不属于 type_set，更新 cur_len为 same_len + 1，并更新 type_set。扫描的过程中记录最大数组长度。

代码

class Solution(object):
    def totalFruit(self, tree):
        """
        :type tree: List[int]
        :rtype: int
        """
        max_len = 1
        cur_len = 1
        same_len = 1
        type_set = set([tree[0]])
        for i in range(1, len(tree)):
            if tree[i] == tree[i-1]:
                cur_len += 1
                same_len += 1
            else:
                if tree[i] in type_set:
                    cur_len += 1
                elif len(type_set) == 1:
                    type_set.add(tree[i])
                    cur_len += 1
                else:
                    type_set = set([tree[i - 1], tree[i]])
                    cur_len = same_len + 1
                same_len = 1
            max_len = max(max_len, cur_len)
        return max_len