Prefix Sum Primes [Codeforces Round #556 (Div. 1) A]

https://codeforces.com/contest/1149/problem/A

1と2からなるN要素の配列Aがある。
これを並び替える操作をする。
操作後に、B[i] = A[0] + A[1] + ... + A[i]である累積和配列Bを作る。
適切に操作をして、配列Bに含まれる素数の数が最大である配列Aを答えよ。

1≦N≦200000

解説

https://codeforces.com/contest/1149/submission/53504399

下の実装はややこしいが、
2 1 2 2 ... 2 2 1 1 ... 1 1
のように、

• 1番目は2
• 2番目は1
• それ以降はまず2で埋める
• 最後に残りを1で埋める

と最適な配列になる。
累積和配列は、2 3から始まるのが適切っぽく、そこから、+2をしていって奇数の素数を洗ってから、
残りの1を使って刻んでいく。
考える方針としては、2以外の素数は奇数なので、奇数をなるべく作っていけば良さそう。

int N, A[201010];
//---------------------------------------------------------------------------------------------------
void _main() {
	cin >> N;
	rep(i, 0, N) cin >> A[i];

	int one = 0, two = 0;
	rep(i, 0, N) {
		if (A[i] == 1) one++;
		else two++;
	}

	vector<int> ans;

	if (one == 0) {
		rep(i, 0, N) ans.push_back(2);
	}
	else if (one == 1) {
		if (two == 0) ans.push_back(1);
		else {
			ans.push_back(2);
			ans.push_back(1);
			rep(i, 0, two - 1) ans.push_back(2);
		}
	}
	else {
		if (two == 0) {
			rep(i, 0, N) ans.push_back(1);
		}
		else {
			ans.push_back(2);
			ans.push_back(1);
			rep(i, 0, two - 1) ans.push_back(2);
			rep(i, 0, one - 1) ans.push_back(1);
		}
	}

	rep(i, 0, N) {
		if (i) printf(" ");
		printf("%d", ans[i]);
	}
	printf("\n");
}