topcoder: ABC

ABC.cpp

#include <string>
#include <cmath>
#include <set>
using namespace std;

class ABC {
	public:
		string createString(int N, int K) {
			if (K == 0) {
				return string(N, 'A');
			}
			int i, j, k;
			int n3 = N/3, r = N%3;
			i = r == 2? (n3+1): n3;
			j = i;
			k = N - i - j;
			if (i*j+j*k+k*i < K) {
				return "";
			}
			j = k = int(N-sqrt(N*N-3*K)+2) / 3;
			if (j == 0) {
				j = 1;
			}
			i = N - j - k;
			string s;
			s.reserve(N);
			for (int p = 0; p < i; p++) { s.push_back('A'); }
			for (int p = 0; p < j; p++) { s.push_back('B'); }
			for (int p = 0; p < k; p++) { s.push_back('C'); }

			set<int> pairs;
			if (i != 0) { pairs.insert(i-1); }
			if (j != 0 && k != 0) { pairs.insert(i+j-1); }

			int curK = i*j+j*k+k*i;
			while (curK > K) {
				set<int>::iterator first = pairs.begin();
				int p = *first;
				pairs.erase(first);
				// swap s[p],s[p+1]
				char tmp = s[p];
				s[p] = s[p+1];
				s[p+1] = tmp;

				if (p - 1 >= 0 && s[p-1] < s[p]) {
					pairs.insert(p-1);
				}
				if (p + 2 < N && s[p+1] < s[p+2]) {
					pairs.insert(p+1);
				}
				curK--;
			}
			return s;
		}
};