問題概要

$N$ 個の正の整数 $A_{1}, A_{2}, \dots, A_{N}$ を $M$ 個のグループに分ける。ただし、どのグループの要素数も 1 個以上 2 個以下でなければならない。

最適なグループ分けをしたときの、各グループの要素の総和の二乗の総和の最小値を求めよ。

制約

$1 \le N \le 2 \times 10^{5}$
$N/2 \le M \le N$

考えたこと

総和が一定なので、二乗和を最小化するためには、なるべく均等にすればよいということがわかる。

よって、 $A$ をソートして、小さい方から $N-M$ 個について、それらを逆順に残りの $M$ 個の $A$ に足していけば良いと予想を立てた。

小さい方から $N - M$ 個を、残りの小さい方から $N - M$ 個とペアを作って行く方法の中に最適解があること
さらに、ペアにする際に大きさの逆順にするのがよいこと

は、それぞれ証明可能である。

計算量は、ソートする部分がボトルネックで $O(N \log N)$ となる。

コード

#include <bits/stdc++.h>
using namespace std;


int main() {
    long long N, M;
    cin >> N >> M;
    vector<long long> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i];
    sort(A.begin(), A.end());
    
    vector<long long> B(M), C(N - M);
    for (int i = 0; i < N - M; ++i) C[i] = A[i];
    for (int i = 0; i < M; ++i) B[i] = A[i + N - M];
    reverse(C.begin(), C.end());
    for (int i = 0; i < N - M; ++i) {
        B[i] += C[i];
    }
    
    long long res = 0;
    for (int i = 0; i < M; ++i) {
        res += B[i] * B[i];
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 167 A - Toasts for Breakfast Party (灰色, 300 点)

問題概要

制約

考えたこと

コード