[LeetCode] 75. Sort Colors

题目链接: https://leetcode.com/problems/sort-colors/description/

Description

Given an array with n objects colored red, white or blue, sort them so that objects of the same color are adjacent, with the colors in the order red, white and blue.

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.

Note:
You are not suppose to use the library’s sort function for this problem.

解题思路

设两个变量 red_next 和 blue_prev，分别代表红色下一个该放的位置和蓝色前一个该放的位置。下一个和前一个是相对数组的位置来说的，数组下标 1 是 2 的前一个，2 是 1 的后一个。

先把数组两端满足顺序的过滤出，只考虑中间不满足的数组元素，过滤后数组下标 red_next 前面的所有元素都为 0，blue_prev 后面的所有元素都为 2。

然后，对数组从 red_next 到 blue_prev 遍历一遍，当发现红色的，即值为 0，就与 nums[red_next] 交换，并将 red_next 加一，这样遍历后所有红色的都放到了数组的前部。

最后，再对数组从 blue_prev 到 red_next 遍历一遍，当发现蓝色的，即值为 2，就与 nums[blue_prev] 交换，并将 blue_prev 减一，这样就将蓝色的都放到了数组的后部。

时间复杂度 O(n)，空间复杂度 O(1)

Code

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int red_next = 0;
        int blue_prev = nums.size() - 1;

        while (red_next < nums.size() && nums[red_next] == 0)
            red_next++;

        while (blue_prev >= 0 && nums[blue_prev] == 2)
            blue_prev--;

        for (int i = red_next; i <= blue_prev; ++i) {
            if (nums[i] == 0) {
                swap(nums[i], nums[red_next]);
                red_next++;
            }
        }

        for (int i = blue_prev; i >= red_next; --i) {
             if (nums[i] == 2) {
                swap(nums[i], nums[blue_prev]);
                blue_prev--;
            }
        }
    }
};