0904-leetcode算法实现之水果成篮-fruit-into-baskets-python&golang实现

在一排树中，第 i 棵树产生 tree[i] 型的水果。
你可以从你选择的任何树开始，然后重复执行以下步骤：

把这棵树上的水果放进你的篮子里。如果你做不到，就停下来。
移动到当前树右侧的下一棵树。如果右边没有树，就停下来。
请注意，在选择一颗树后，你没有任何选择：你必须执行步骤 1，然后执行步骤 2，然后返回步骤 1，然后执行步骤 2，依此类推，直至停止。

你有两个篮子，每个篮子可以携带任何数量的水果，但你希望每个篮子只携带一种类型的水果。

用这个程序你能收集的水果树的最大总量是多少？

示例 1：

输入：[1,2,1]
输出：3
解释：我们可以收集 [1,2,1]。
示例 2：

输入：[0,1,2,2]
输出：3
解释：我们可以收集 [1,2,2]
如果我们从第一棵树开始，我们将只能收集到 [0, 1]。
示例 3：

输入：[1,2,3,2,2]
输出：4
解释：我们可以收集 [2,3,2,2]
如果我们从第一棵树开始，我们将只能收集到 [1, 2]。
示例 4：

输入：[3,3,3,1,2,1,1,2,3,3,4]
输出：5
解释：我们可以收集 [1,2,1,1,2]
如果我们从第一棵树或第八棵树开始，我们将只能收集到 4 棵水果树。

提示：

1 <= tree.length <= 40000
0 <= tree[i] < tree.length

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/fruit-into-baskets

python

# 水果成篮
class Solution:
    def totalFruit(self, tree: [int]) -> int:
        """
        滑动窗口+双指针+哈希表， 时间O(n), 空间O(1)
        我们来抽象一下，就是给你一个数组， 让你选定一个子数组，
        这个子数组最多只有两种数字，这个选定的子数组最大可以是多少。
        思路：
        借助辅助哈希存储窗口元素出现次数，滑动窗口的左右边界不断调整使得窗口内仅含2种不同元素值，一次遍历，即可得到滑动窗口的最大值
        详细流程：
        1.初始化窗口长度ans为0，左边界指针i指向0号元素，初始化哈希表win，对应值初始0
        2.从第一个元素开始遍历，作为滑动窗口的右边界指针j,不断移动右边界，即滑动窗口不懂右移的过程，j++
        3.检查当前右边界指针元素是否出现过，窗口中未含该元素，窗口可容纳元素种类k--
        4.不管右边界指针元素是否出现过，其在窗口中的出现次数++
        5.检查当前边界窗口的可容纳元素种数是否小于0(-1)，如果小于0，即当前窗口含3种元素，需要不断移动左边界指针，使得最终窗口内元素种类小于等于2，即k>=0
            -上述检查通过，即左边界指针元素出现次数--
            -左边界指针元素出现次数需要检查是否减完，即窗口哈希win中对应值为0，此时意味窗口可容纳元素种类数++
            -向右移动左边界指针i++
        6.上述3/4/5步骤走完，即当前窗口的左边界指针已停止移动，符合题目元素种类数要求，求得当前窗口覆盖的子数组的长度j-i+1，与上轮或初始化的ans比较得到较大值作为新的ans
        7.重复2/3/4/5/6步骤，直到所有元素都遍历完，最后的ans即为题中要求的窗口长度最大值
        :param tree:
        :return:
        """
        def atMostK(k, nums):
            """
            滑动窗口取k个不同值元素的最长连续子数组长度
            :param k:
            :param nums:
            :return:
            """
            from collections import defaultdict
            i = ans = 0 # 初始化左指针或左边界，暂且叫左边界，初始化返回结果为0，即仅含2种元素的最长子数组长度
            win = defaultdict(lambda: 0) # 初始化哈希表，val=0，val表示出现次数
            for j in range(len(nums)): # 仅需遍历一轮数组
                if win[nums[j]] == 0: # 每次遍历数组元素时，检查在哈希表中是否存在该元素，如果存在，必定val已经大于0，不存在，则必定是初始化为0的或者减为0
                    k -= 1 # 如果元素对应出现次数为0，意味窗口中加入了新值元素，对应窗口含有的元素种类减1
                win[nums[j]] += 1 # 每次遍历元素的时候，都在对应的val上+1
                while k < 0: # 遍历每个元素时，检查k是否小于0，说明窗口中可定含有2种以上不同值的元素(3种不同值的元素)，此时应该考虑改变窗口边界，主要是移动窗口左边界，最终应该使得窗口内含有2种不同值元素，即k=0
                    win[nums[i]] -= 1 # 每次while循环，当前左边界指针元素对应的次数减1
                    if win[nums[i]] == 0: # 检查左边界指针对应元素出现次数是否为0，意味着3种不同值变为了2种，此时k由-1变为0，即k++
                        k += 1
                    i += 1 # 检查左边界指针对应次数不为0，意味着左边界应该右移，即i++，移动左边界指针时检查元素种类是否大于2(3)，只要元素值种数为3，就移动左边界，直到k=0为止
                ans = max(ans, j - i + 1) # 当前符合2种元素值的左右边界指针，取子数组长度及上轮窗口大值的较大值，作为当前遍历元素的大值
            return ans # 当所有元素遍历完，ans即为符合条件滑动窗口的最大值

        return atMostK(2, tree) # 水果成篮，给定一排树，最多含有2种水果的情况

    def totalFruit1(self, tree: [int]) -> int:
        """
        按块扫描，暴力法
        tip:
            0<tree[i]<len(tree)
            1<=len(tree)<=4000
        思路：

        :param tree:
        :return:
        """
        res = 0
        n = len(tree)
        for i in range(n): # 1层循环
            pre1, pre2 = -1, -1 # 每轮循环时，初始化-1
            for j in range(i, n): # 进入2层for loop
                if pre1 == -1 or pre1 == tree[j]: # 如果pre1=-1，表示1层for loop重新开始遍历，或者当pre1等于当前元素的值(意味着遇到与前2个元素相同值的元素)，进入if语句，赋值，continue,j++
                    pre1 = tree[j]
                    continue
                elif pre2 == -1 or pre2 == tree[j]: # 如果pre2=-1，表示1层for loop重新开始遍历，或者当pre2等于当前元素的值(意味着遇到与前2个元素相同值的元素)，进入if语句，赋值，continue,j++
                    pre2 = tree[j]
                    continue
                else: # 如果当前元素不等于pre1或pre2，表示遇到了第三个不同元素的值，从水果成篮的角度，就有第三种水果，意味当前的2层for loop要break，break前获取到连续子数组长度大值
                    if tree[j] == pre1 or tree[j] == pre2:
                        continue
                    res = max(res, j-i) # 由于此时的j为第三个不同元素的值，直接j-i即为当前唯2元素的最长子数组长度
                    break # break后，开始下一轮循环

        return res

    def totalFruit2(self, tree: [int]) -> int:
        """
        暴力扫描，力扣超时间。时间O(n^2), 空间O(3n)->n轮，每轮3对k-v
        思路：
        分2层遍历，第一层遍历的元素作为符合条件的下界，第二层遍历的元素作为上界，借助哈希表存储不同值的元素，
        第二层遍历时，只要哈希表遇到第三种元素，break，每次新的下界更新时，哈希表初始化
        算法流程：
        1.数组长n,初始化元素种类k,初始化返回结果res 0
        2.一层便利遍历i，初始化哈希表
        3.二层遍历j，从i到n-1
        -对j对应元素出现次数++
        -检查哈希表长度是否大于K,是则记录此轮res max,继续下轮遍历
        -对于j已经是最后的元素，依然只有k种元素，记录此轮res max
        4.重复1、2、3步骤，返回res值
        :param tree:
        :return:
        """
        from collections import defaultdict
        n = len(tree)
        k = 2
        res = 0
        for i in range(n):
            subarray_map = defaultdict(int)
            for j in range(i, n):
                subarray_map[tree[j]] += 1
                if len(subarray_map) > k:
                    res = max(res, j-i)
                    break
                if j == n-1:
                    res = max(res, j-i+1)
        return res




if __name__ == "__main__":
    tree = [1,2,1,0,1,1,3,3]
    tree1 = [1,2,1]
    test = Solution()
    print(test.totalFruit2(tree))

golang

package main

import (
	"fmt"
)

func main() {
	nums := []int{1, 2, 1}
	fmt.Println(totalFruit(nums))
}

// 滑动窗口
func totalFruit(tree []int) int {
	return atMostK(2, tree)
}

func atMostK(k int, nums []int) int {
	i := 0
	ans := 0
	win := make(map[int]int)
	for j := 0; j < len(nums); j++ {
		if win[nums[j]] == 0 {
			k--
		}
		win[nums[j]]++
		for k < 0 {
			win[nums[i]]--
			if win[nums[i]] == 0 {
				k++
			}
			i++
		}
		if ans > j-i+1 {
			ans = ans
		} else {
			ans = j - i + 1
		}
	}
	return ans
}