三色排序问题/(荷兰国旗问题)(C++版)

题目：

有一个只由0，1，2三种元素构成的整数数组，请使用交换、原地排序而不是使用计数进行排序。

给定一个只含0，1，2的整数数组A及它的大小，请返回排序后的数组。

算法复杂度：

时间 O(N)

空间 O(1)

贴代码：

void swap(int *p1, int *p2)
	{
		int temp = *p1;
		*p1 = *p2;
		*p2 = temp;
	}

	vector<int> sortThreeColor(vector<int> A, int n) 
	{
		// 0 区域 2区域
		int idx0 = 0, idx2 = n - 1;

		for(int i = 0; i < n;)
		{
			// ==0 交换当前元素和0区的下一个位置(右) 遍历下一个
			if(A[i] == 0)
			{
				swap(&A[i], &A[idx0]);
				idx0++;
				++i;
			}
			// ==2 交换当前元素和2区的下一个位置(左)
			else if(A[i] == 2)
			{
				// 还未遍历到2区 交换
				if (i < idx2)
				{
					swap(&A[i], &A[idx2]);
					idx2--;
				}
				// 否则 直接遍历下一个
				else
				{
					++i;
				}
			}
			// ==1 遍历下一个
			else if(A[i] == 1)
			{
				++i;
			}
		}

		return A;
	}

测试函数：

int main(void)
{
	int a[6] = {0,1,1,0,2,2};

	vector<int> v(a, a + 6);

	ThreeColor tc;

	v= tc.sortThreeColor(v, 6);

	for (int i = 0; i < 6; ++i)
	{
		cout<<v[i];
	}

	return 0;
}

输出结果：
001122