LeetCode in Kotlin

904. Fruit Into Baskets

Medium

You are visiting a farm that has a single row of fruit trees arranged from left to right. The trees are represented by an integer array fruits where fruits[i] is the type of fruit the ith tree produces.

You want to collect as much fruit as possible. However, the owner has some strict rules that you must follow:

Given the integer array fruits, return the maximum number of fruits you can pick.

Example 1:

Input: fruits = [1,2,1]

Output: 3

Explanation: We can pick from all 3 trees.

Example 2:

Input: fruits = [0,1,2,2]

Output: 3

Explanation: We can pick from trees [1,2,2]. If we had started at the first tree, we would only pick from trees [0,1].

Example 3:

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

Output: 4

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

Constraints:

Solution

class Solution {
    fun totalFruit(fruits: IntArray): Int {
        var end = 1
        var basket1 = fruits[0]
        var basket2 = -1
        var secondFruitIndex = -1
        var maxTotal = 1
        var counter = 1
        while (end < fruits.size) {
            if (fruits[end - 1] != fruits[end]) {
                if (basket2 == -1) {
                    basket2 = fruits[end]
                    secondFruitIndex = end
                    counter++
                } else if (fruits[end] == basket1) {
                    basket1 = basket2
                    basket2 = fruits[end]
                    secondFruitIndex = end
                    counter++
                } else {
                    counter = end - secondFruitIndex + 1
                    basket1 = basket2
                    basket2 = fruits[end]
                    secondFruitIndex = end
                }
            } else {
                counter++
            }
            end++
            maxTotal = maxTotal.coerceAtLeast(counter)
        }
        return maxTotal
    }
}