Remove Duplicates from Sorted Array II

class Solution:
    def removeDuplicates(self, nums: List[int]) -> int:
        k = 0
        count = 0

        for i in range(len(nums)):

            # New number, reset count
            if i == 0 or nums[i] != nums[i - 1]:
                count = 1
            else:
                count += 1

            # Keep at most 2 occurrences
            if count <= 2:
                nums[k] = nums[i]
                k += 1

        return k