An orchard has one long row of trees, and trees[i] is the kind of fruit that tree i grows (kinds are numbered). You carry two baskets. Each basket holds as much fruit as you like, but only of one kind.
You pick a starting tree, then walk right along the row, picking exactly one fruit from every tree you pass, starting tree included. You can't skip a tree. As soon as you reach a tree whose kind matches neither basket (when both already hold something), you stop. You may also stop at the end of the row.
Write total_fruit(trees: list[int]) -> int: the most fruit you can collect.
total_fruit([4, 4, 7, 4, 9, 9, 7, 9, 9, 4]) # 5: start at tree 4 and pick 9, 9, 7, 9, 9
total_fruit([5, 2, 5, 8]) # 3: start at tree 0 and pick 5, 2, 5
total_fruit([3, 3, 3]) # 3: one basket is enough
total_fruit([1, 2, 3, 4, 5]) # 2: any two neighbours
Constraints: 1 <= len(trees) <= 2 * 10^5, 0 <= trees[i] < len(trees).
⭐ Bonus: walking right from every possible start still passes; a speed test with 200,000 trees earns a star in O(n).
Show hint
you're looking for the longest stretch of the row that holds at most two different kinds. Keep a count of each kind in the current stretch, and when a third kind comes in, drop trees from the left until one kind is gone.