76 leetcode - Remove Duplicates from Sorted Array II

#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
Remove Duplicates from Sorted Array II
Follow up for "Remove Duplicates":
What if duplicates are allowed at most twice?

For example,
Given sorted array nums = [1,1,1,2,2,3],

Your function should return length = 5, with the first five elements of nums being 1, 1, 2, 2 and 3. 
It doesn't matter what you leave beyond the new length.
'''
class Solution(object):
    def removeDuplicates(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        length = len(nums)
        if length <= 2:
            return length

        flag = False
        ret = 1
        cur = 1
        for index in range(1,length):
            if nums[index] == nums[index - 1]:
                if flag == False:
                    ret += 1    
                    flag = True
                    nums[cur] = nums[index]
                    cur += 1
            else:
                flag = False
                ret += 1
                nums[cur] = nums[index]
                cur += 1
        #print nums[:cur]
        return ret

if __name__ == "__main__":
    s = Solution()
    print s.removeDuplicates([1,1,1,2,2,3])
    print s.removeDuplicates([1,1,1,1,1,1,2,2,2,3])