leetcode/75.Sort Colors/75.cpp at master · arkingc/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution {
public:
    void sortColors(vector<int>& nums) {
        int redcount = 0, whitecount = 0;

        for(int i = 0;i < nums.size();i++){
            if (nums[i] == 0)
                redcount++;
            else if(nums[i] == 1)
                whitecount++;
            else if(nums[i] == 2)
                continue;
            else
                return;
        }

        for(int i = 0;i<nums.size();i++){
            if(redcount){
                nums[i] = 0;
                redcount--;
            }
            else if(whitecount){
                nums[i] = 1;
                whitecount--;
            }
            else
                nums[i] = 2;
        }
    }
};