LeetCode C++18-颜色分类

题目描述

        给定一个包含红色、白色和蓝色、共n个元素的数组nums，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。

        我们用整数0、1和2分别表示红色、白色和蓝色。注意：必须在不使用库的sort函数的情况下解决这个问题。

举例

示例1：输入: nums=[2,0,2,1,1,0]    输出: [0,0,1,1,2,2]

示例2：输入: nums=[2,0,1]             输出：[0,1,2]

示例3:  输入: nums=[2,0,1,2,1,2]    输出:[0,1,1,2,2,2]

解题方法： 双指针，left指针指向左边连续0的下一个位置，right指针指向右边连续2的前一个位置。

代码

#include <iostream>
#include <vector>

void sort_colors(std::vector<int>& nums) {
	int n = nums.size();
	if (n <= 0) {
		return;
	}
	int p0 = 0; // 始终指向左侧连续0的下一个位置
	int p2 = n - 1; //始终指向右侧连续1的前一个位置
    int i = 0;
	while (i <= p2) {
		if (nums[i] == 2) {
			std::swap(nums[i], nums[p2]);
			p2--;//注意：i不能++，因为如果i++，原来p2指向2的话，就会被跳过
		} else if (nums[i] == 0) {
			//[p0, i)之间的元素比如不会有0，如果有0，
			// 则在i扫描的时候，肯定会被swap到左侧
			std::swap(nums[i], nums[p0]);
			p0++;
			i++;
		} else {//i指向的元素为1
			i++;
		}
	}
}

void print_arr(const std::vector<int>& input) {
	for (auto value : input) {
		std::cout << value << " ";
	}
	std::cout << std::endl;
}

int main() {
	// case1: nums=[2,0,2,1,1,0] ==> [0,0,1,1,2,2]
	std::vector<int> nums = {2, 0, 2, 1, 1, 0};
    std::cout << "input:" << std::endl;
    print_arr(nums);
    std::cout << "output:" << std::endl;
    sort_colors(nums);
    print_arr(nums);
	// case2: nums=[2,0,1] ==> [0,1,2]
	nums = {2, 0, 1};
    std::cout << "input:" << std::endl;
    print_arr(nums);
    std::cout << "output:" << std::endl;
    sort_colors(nums);
    print_arr(nums);
    return 0;
}

代码运行结果如下：

input:
2 0 2 1 1 0 
output:
0 0 1 1 2 2 
input:
2 0 1 
output:
0 1 2