USACO Section 2.1 Sorting a Three-Valued Sequence

/*
ID: lucien23
PROG: sort3
LANG: C++
*/

#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>
using namespace std;
void exchange(int nums[], int begin, int end, int N, int x);
int sum = 0;
int main()
{
	ifstream infile("sort3.in");
	ofstream outfile("sort3.out");
	if(!infile || !outfile)
	{
		cout << "file operation failure!" << endl;
		return -1;
	}

	int N;
	infile >> N;
	int *nums = new int[N];
	int count1, count2, count3;
	count1 = count2 = count3 = 0;
	for (int i=0; i<N; i++)
	{
		infile >> nums[i];
		switch (nums[i])
		{
		case 1:
			count1++;
			break;
		case 2:
			count2++;
			break;
		case 3:
			count3++;
			break;
		default:
			break;
		}
	}

	exchange(nums, 0, count1, N, 1);
	exchange(nums, count1, count1+count2, N, 2);

	outfile << sum << endl;

	return 0;
}

/*
 *交换时将大的交换到后面，小的交换到前面
 */
void exchange(int nums[], int begin, int end, int N, int x)
{
	int count0 = 0;
	vector<int> vecNum;
	for (int i=begin; i<end; i++)
	{
		if (nums[i] != x)
		{
			count0++;
			vecNum.push_back(nums[i]);
			nums[i] = x;
		}
	}
	sum += count0;

	sort(vecNum.begin(), vecNum.end());
	for (int i=end, j=0; i<N && count0>0; i++)
	{
		if (nums[i] == 1)
		{
			nums[i] = vecNum[j];
			j++;
			count0--;
		}
	}
}