Leetcode 75. Sort Colors (python+cpp)

题目

解法1：two pass

first pass统计0，1，2出现的次数，second pass对不同的范围内进行重新赋值，这样需要额外的空间储存三种颜色出现的频率，代码如下：

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        
        memo= collections.defaultdict(int)
        for num in nums:
            memo[num] += 1
        
        
        count0 = memo[0]
        count1 = memo[1]
        count2 = memo[2]
        
        for i in range(len(nums)):
            if i<count0:
                nums[i] = 0
            elif count0<=i<count0+count1:
                nums[i] = 1
            else:
                nums[i] = 2

C++版本

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int red = 0, white = 0, blue = 0;
        for(int num : nums){
            if(num == 0){
                red++;
            }else if(num == 1){
                white++;
            }else{
                blue++;
            }
        }
        for(int i=0;i<nums.size();i++){
            if(i < red){
                nums[i] = 0;
            }else if(i >= red && i < red + white){
                nums[i] = 1;
            }else{
                nums[i] = 2;
            }
        }
    }
};

解法2：三指针法

核心思想类似于quicksort的过程，在遍历过程中，想办法将0，2分别放到属于他们的位置区间，具体步骤如下：

• 初始化p0,p1和curr。p0代表0元素现在应该放置的位置，p1代表2元素现在应该放置的位置，curr代表现在正在遍历的元素
• 如果curr现在所指元素为0，那么和p0所指位置元素交换，并且curr和p0都加1
• 如果curr现在所指元素为2，则与p1所指元素交换，并且p1减一，但是curr保持不变
• 终止条件为curr所指位置大于p1所指位置

下面首先来分析一下终止条件，为什么终止条件是curr>p1呢，因为我们是从前往后遍历的，所以curr之前的元素都已经交换完毕，也就是说curr之前的元素都在他们该在的位置，而p1之后的元素毫无疑问都应该是2，也都在他们该在的位置，所以一旦curr位置超过p1位置，就证明整个数组都已经操作完毕

其次分析为什么在遇到0并进行交换后，curr的位置需要后移，而遇到2时则不能前移。举个很简单的例子，如果在遇到2后curr也后移，那么对于[1,2,0]这样的情况，我们得到的结果会是[1,0,2]。为什么会出现这种情况呢？这个根本原因是因为我们是从前往后遍历的，所以当遇到0时，我们跟p0位置的元素交换，换过来的只可能是0或者1，也就是说对curr交换后的这个位置元素无需进行进一步操作，而对2的情况，我们并不能确定换到curr的元素是否需要进一步操作，如果跟2交换的元素是0，那么意味着curr这个位置的0可能需要进一步操作。

根据上面的分析，如果我们从后往前遍历，那么curr就应该在遇到0时，交换后需要进一步判断，而遇到2时不需要。下面提供了两种不同的遍历方向的解法

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        # 从后往前遍历，遇0交换后需要进一步判断
        p0 = 0
        p2 = curr = len(nums)-1
        
        while curr>=p0:
            if nums[curr]==0:
                nums[curr],nums[p0] = nums[p0],nums[curr]
                p0 += 1
            elif nums[curr]==2:
                nums[curr],nums[p2] = nums[p2],nums[curr]
                p2 -= 1
                curr -= 1
            else:
                curr -= 1
        
        
        # 从前往后遍历，遇2交换需要进一步判断
        p0 = curr = 0
        p2 = len(nums)-1
        
        while curr<=p2:
            if nums[curr]==0:
                nums[curr],nums[p0] = nums[p0],nums[curr]
                curr += 1
                p0 += 1
            elif nums[curr]==2:
                nums[curr],nums[p2] = nums[p2],nums[curr]
                p2 -= 1
            else:
                curr += 1

C++版本

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int p0 = 0, p2 = nums.size() - 1,pcurr = 0;
        while(pcurr <= p2){
            if(nums[pcurr] == 0){
                // increment p0 and pcurr both
                // because the nums[p0] before swap can only be 1 or 0
                // if nums[p0] before swap, it will be already swapped previously
                swap(nums[p0],nums[pcurr]);
                p0++;
                pcurr++;
            }else if(nums[pcurr] == 2){
                // only decrement p2 because we can not sure if pcurr is at right position
                swap(nums[pcurr],nums[p2]);
                p2--;
            }else{
                // nums[pcurr] == 1, it is already at the right position
                pcurr++;
            }
        }
    }
};