数列を2つの等差数列に分割する

COCI 2019 #3 Drvca

問題

入力で整数列Aが与えられる。それを2つの数列P, Qに分割する。順番は並べ替えてもいい。P, Qを等差数列にするようにできるか？

解法

Aのサイズが4以下なら半分に分けて自明。
それより多い場合。Aはソートされているとする。
Aの小さい3つの値を小さい順にx, y, zとする。P,Qの少なくとも一方はx, y, zの少なくとも2つを含むので、x, y, zから2つ選ぶと初項と公差が決定する。
それをPとして、Pを1つずつ伸ばしたときに、残りの要素が等差数列になっているか判定できていればいい。
これは「残りの要素」をmultisetに、「残りの要素の隣り合う値の差」をmultisetにして更新し、multisetが空または1種になればよい。

vector<int> A = { 1, 5, 8, 9, 10, 12 };
// P = 1, 5, 9;
// Q = 8, 10, 12;
vector<int> P;
vector<int> Q;
bool solve() {
    if (A.size() <= 4) {
	REP (i, A.size()) {
	    if (i % 2 == 0) P.push_back(A[i]);
	    else Q.push_back(A[i]);
	}
	return true;
    } else {
	sort(A.begin(), A.end());
	for (int i=0; i<2; i++) for (int j=i+1; j<3; j++) {
	    int x = A[i];
	    int d = A[j] - x;
	    multiset<int> elems(A.begin(), A.end());
	    multiset<int> diffs;
	    REP (k, (int)A.size()-1) diffs.insert(A[k+1] - A[k]);
	    while (1) {
		if (diffs.empty() || *diffs.begin() == *diffs.rbegin()) {
		    Q.assign(elems.begin(), elems.end());
		    return true;
		}
		auto it = elems.find(x);
		if (it == elems.end()) {
		    P.clear();
		    break;
		}
		auto right = next(it);
		if (right != elems.end()) diffs.erase(diffs.find(*right - *it));
		if (it != elems.begin()) {
		    auto left = prev(it);
		    diffs.erase(diffs.find(*it - *left));
		    if (right != elems.end()) diffs.insert(*right - *left);
		}
		elems.erase(it);
		P.push_back(x);
		x += d;
	    }
	}
	return false;
    }
}