[Array] LeetCode 80. Remove Duplicates from Sorted Array II

Solution O(n)

class Solution:
    def removeDuplicates(self, nums: List[int]) -> int:
        n = len(nums)
        current_num = nums[0]
        count, idx = 1, 0

        for i in range(1,n):
            if nums[i] == current_num:
                count += 1

            if nums[i] != current_num:
                if count >= 3:
                    count = 2

                for _ in range(count):
                    nums[idx] = current_num
                    idx += 1

                count = 1
                current_num = nums[i]

        if count >= 3:
            count = 2

        for _ in range(count):
            nums[idx] = current_num
            idx += 1

        while idx < n:
            nums.pop()
            idx += 1

1. If nums[i] == current_num, then we add 1 to count.
2. If nums[i] != current_num, then we first check if count is bigger than 3. If it is bigger than 3, we change count to 2. After that, we update the nums array with current_num.
3. Reset count to 1 and current_num to the current number nums[i].
4. After iterating through the whole array, we have to check once again to update the nums array for the last element.
5. We pop the last elements of the nums array if there are leftover numbers.