LeetCode() Sort Colors

一个数组里有0,1,2三种数，排列

非常牛逼的思路！！

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int i=-1,j=-1,k=-1;
        
        for(int p=0;p<nums.size();++p)
        {
            if(nums[p] == 0)
            {
                nums[++k]=2;
                nums[++j]=1;
                nums[++i]=0;
                
            }
            else if(nums[p] == 1)
            {
                nums[++k]=2;
                nums[++j]=1;
            }
            else
                nums[++k]=2;
        }
    }
};

给跪了，如果有4种不同的数，应该也可以这么做，这个思路真好。为什么我的有问题呢？

入下

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int start=0,end=nums.size()-1;
        while(start<end)
        {
            while(nums[start] != 2)
                start++;
            while(nums[end] == 2)
                end--;
            if(start<end)
                swap(nums[start],nums[end]);
        }
        start=0;
        while(start<end)
        {
            while(nums[start] != 1)
                start++;
            while(nums[end] == 1)
                end--;
            if(start<end)
                swap(nums[start],nums[end]);
        } 
    }
};