Collect maximum fruits in two baskets of at most 2 types of fruits The task

is taken from LeetCode


  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:
  
  
  Add one piece of fruit from this tree to your baskets.  If you cannot, stop.
  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 &lt;= tree.length &lt;= 40000
  0 &lt;= tree[i] &lt; tree.length
  


My solution:

has time and space complexity of \$O(n)\$. At first I thought it was easy. But then, I got confused with the one of the test cases (i.e. I = [3,3,3,1,2,1,1,2,3,3,4];) and everything inside the else-block is a bit hacky afterwards. Maybe there is a more elegant solution to that.

/**
 * @param {number[]} tree
 * @return {number}
 */
var totalFruit = function(tree) {
    const set = new Set(tree);
    if (set.size &lt;= 2) {
        return tree.length;        
    }

    const fruits = new Set();
    let i = 0;
    let j = 0;
    let max = 0;
    let count = 0;
    while (j &lt; tree.length) {
        if (fruits.size &lt;= 2 &amp;&amp; !fruits.has(tree[j])) {
            fruits.add(tree[j]);
        }

        if (fruits.size &lt;= 2) {
            count++;
            max = Math.max(max, count);
            j++;
        } else {
            fruits.delete(tree[i]);
            const lastIndex = tree.slice(i, j - 1).lastIndexOf(tree[i]);
            i += lastIndex + 1;
            count-= lastIndex + 1;
        }
    }
    return max;
};

let I = [1,2,1];
I = [0,1,2,2];
I = [3,3,3,1,2,1,1,2,3,3,4];
console.log(totalFruit(I));


Update: I made a mistake. This should be the accurate solution:

/**
 * @param {number[]} tree
 * @return {number}
 */
var totalFruit = function(tree) {
    let max = 0, count = 0;
    for (let i = 0, first = 0, second = -1; i &lt; tree.length; i++) {
        count++;
        if (tree[i] === tree[first]) {
            first = i;
        } else if (second === -1 || tree[i] === tree[second]) {
            second = i;
        } else {
            max = Math.max(count - 1, max);
            count = Math.abs(first - second) + 1;
            first = i - 1;
            second = i;
        }
    }
    return Math.max(count, max);
};


Time complexity \$O(n)\$ and space complexity \$O(1)\$
Your second solution is good and fast, but you asked about more elegant solution, so I am suggesting mine. Firstly, I solved this task by Python, then convert all logic into Javascript. It is a little slower, than your (20 ms) and use more memory, but I think it more straightforward and understandable.

The code only answers do not liked on this site, so I add some comparisons:


Algorithm. Both algorithms are similar but:


Mine: keeps track of every number changing positions. The start1 position changes every time the number was changed, so I always know the position, where the previous number was started. The start2 position changes only when third number occurs, so I just subtract the start2 from the current index and get the needed two number sequence length.
Your: keeps track of last occurrences the first and second numbers, so you miss their start positions, and thus, you should use the count variable for storing the length of the current two number sequence. When the third number appears, you need to calculate the value of the last uninterruptible one number sequence by Math.abs(first - second). Also, you don't know which number was last - first or second, so the Math.abs function is needed.

The way of access to array items. 


Mine: uses the iterator - for...of statement and Array.entries(). It relieves us from tree[i] and tree[second] like constructions.
Your: Uses counter and array indexes to access the needed item.



The Javascript code:

var totalFruit = function(tree) {
    let n1 = -1;
    let n2 = -1;

    let start1 = 0;
    let start2 = 0;

    let maxim = 1;
    # Add extra element in the end of array (which is not occured in array)
    # to get rid of the second 'Math.max(maxim, k - start2)' call
    tree.push(-2);

    for (let [k, num] of tree.entries()) {
        if (num !== n1) {
            if (num !== n2) {
                maxim = Math.max(maxim, k - start2);
                start2 = start1;
            }

            n2 = n1;
            n1 = num;

            start1 = k;
        }
    }

    return maxim;
}


The original Python code:

class Solution:
    def totalFruit(self, tree):
        n1 = -1
        n2 = -1

        start1 = 0
        start2 = 0

        maxim = 1
        tree.append(-2)

        for k, num in enumerate(tree):
            if num != n1:
                if num != n2:
                    maxim = max(maxim, k - start2)

                    start2 = start1

                n2 = n1
                n1 = num

                start1 = k

        return maxim