Formatted question description: https://leetcode.ca/all/904.html

Medium

## Description

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:

2. 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].

Example 3:

Input: [1,2,3,2,2]

Output: 4

Explanation: We can collect [2,3,2,2]. If we started at the first tree, we would only collect [1, 2].

Example 4:

Input: [3,3,3,1,2,1,1,2,3,3,4]

Output: 5

Explanation: We can collect [1,2,1,1,2]. If we started at the first tree or the eighth tree, we would only collect 4 fruits.

Note:

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

## Solution

This problem is equivalent to the following problem: Given an array, find the length of the longest subarray such that there are at most 2 distinct numbers in the subarray.

Use a map to store each number and its number of occurrences. Loop over the array tree once from left to right, and maintain a variable start that stores the start index of the current subarray. For each number, increase its number of occurrences by 1 in the map. While the map contains more than 2 keys, move start forward and reduce the number of occurrences of the integers that start pointed to previously. If a number’s number of occurrences becomes 0, then remove the number from the map. Repeat the process until the map contains at most 2 keys. Then update the length of the longest subarray that contains at most 2 distinct numbers, which equals the current index minus start. Finally, return the maximum length.

• class Solution {
public int totalFruit(int[] tree) {
Map<Integer, Integer> typeCountMap = new HashMap<Integer, Integer>();
int maxTotalCount = 0;
int start = 0;
int length = tree.length;
for (int i = 0; i < length; i++) {
int type = tree[i];
int count = typeCountMap.getOrDefault(type, 0);
count++;
typeCountMap.put(type, count);
while (typeCountMap.keySet().size() > 2) {
int prevType = tree[start];
int prevCount = typeCountMap.getOrDefault(prevType, 0);
if (prevCount > 0)
prevCount--;
if (prevCount == 0)
typeCountMap.remove(prevType);
else
typeCountMap.put(prevType, prevCount);
start++;
}
maxTotalCount = Math.max(maxTotalCount, i - start + 1);
}
return maxTotalCount;
}
}

• // OJ: https://leetcode.com/problems/fruit-into-baskets/
// Time: O(N)
// Space: O(1)
class Solution {
public:
int totalFruit(vector<int>& A) {
int i = 0, j = 0, N = A.size(), ans = 0;
unordered_map<int, int> m;
for (; j < N; ++j) {
m[A[j]]++;
while (m.size() > 2) {
if (--m[A[i]] == 0) m.erase(A[i]);
++i;
}
ans = max(ans, j - i + 1);
}
return ans;
}
};

• class Solution(object):
def totalFruit(self, tree):
"""
:type tree: List[int]
:rtype: int
"""
left, right = 0, 0
res = 0
cnt = collections.defaultdict(int)
while right < len(tree):
cnt[tree[right]] += 1
while len(cnt) > 2:
cnt[tree[left]] -= 1
if cnt[tree[left]] == 0:
del cnt[tree[left]]
left += 1
res = max(res, right - left + 1)
right += 1
return res