【AtCoder】ABC136E Max GCD

問題

https://atcoder.jp/contests/abc136/tasks/abc136_e
どれか2つの要素を選んで一方を+1、もう一方を-1するという操作を0回以上K回以下行い、全要素に対する最大公約数の最大値を求める問題。

模範解答では累積和を使っていたが、使わなくて良い方法があり、そのことを記事にしておりかつソースコードを公開している人が見当たらなかったので書く。

解法

まず操作による各a_1~a_nの変化量をd_1~d_nとおくと、dの和は0になる。(操作の仕組みから。)
最大公約数をgとおくと、a_1+d_1~a_n+d_nは全部gの倍数だからa+dの総和もgの倍数。dの総和は0なのでaの総和がgの倍数になる。
よって答えの候補をまず総和の約数に絞る。(ここまではすぐだった……)

次にこの候補gに対してK回以内の操作で実現可能かを考える。
a_1~a_nを全てgで割った余りr_1~r_nに置き換える。(rは0以上g未満)
このときr_1~r_nは操作によって0またはgになる。(-gや2gや3g…にはならない)(まあさすがに無駄っぽく見えるけど……)

お気持ち証明：あるr_iについて、操作の結果、2g以上(2g,3g,4g…)になったとき、d_i≧g+1であり、dの総和が0だから、Σd-d_i=-(g+1)<0
よって、d_j<0となるjが存在する。このようなd_iとd_jに対して、d_iをd_i-g、d_jをd_j+gを置き換えても、dの総和は0で変わらない。
操作回数について。+方向の操作回数について、まずr_iに対してg回操作回数が減る。
次にr_jに対しては、d_jがもともと負であったことからd_j+gが負の時は+方向の操作回数には影響がなく、もしd_j+gが正になってもd_j+gはg未満なのでg回以上操作回数が増えることはない。結果として+方向の操作回数は真に減るので操作回数自体も減り、ゆえに2g以上にする意味はない。
同様にして-g以下の場合も示す。(多分できるはず)

ということでr_1~r_nを0またはgにするので任意のiについてd_iは-r_iまたはg-r_iになる。
これをうまく組み合わせてdの総和が0、かつ操作回数を最小にしたときにK回以下かがわかれば良い。
まず全部d_i=-r_iにしてみる。(r_iを全部0にする)
するとΣrがgの倍数なので(これはΣaがgの倍数であることからわかる)、Σd=-Σrとなり、Σdはgの倍数にはなる。が0とは限らない。(というか、d_iが全部負なので、0にならない)
ここで-r_iをg-r_iにするとΣdはg増えているので、Σdを0にするためには、(-Σd)/g個の項d_iを-r_iからg-r_iにすれば良い。
どれを-r_iからg-r_iに変えると操作回数が減りそうか？今-方向の操作回数はΣr回で+方向の操作回数は0回なので、rが大きい順に-r_iからg-r_iに変えれば良いことが分かる。
ということで大きい(-Σd)/g=(Σr)/g項を取り除くと-方向の操作回数になるのでこれがK回以下かを判断すれば良い。

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

typedef long long int ll;
#define rep(i,n) for (int i = 0; i < (n); i++)

vector<int> divisor(int N) {
    vector<int> RES;
    for (int I = 1; I*I <= N; I++) {
        if (N%I == 0) {
            RES.push_back(I);
            if (I != N/I) RES.push_back(N/I);
        }
    }
    return RES;
}

int main() {
    int n; cin >> n;
    ll k; cin >> k;
    int a[n]; rep(i,n) cin >> a[i];

    int s = 0;
    rep(i,n) s += a[i];
    vector<int> d = divisor(s);

    int ans = 1;

    for (auto g : d) {
        vector<int> r;
        rep(i,n) r.push_back(a[i]%g);
        sort(r.begin(),r.end());

        int sum_r = 0; rep(i,n) sum_r += r[i];
        int cmax = sum_r / g;
        int c = 0;
        rep(i,n-cmax) c += r[i];
        if ((ll)c <= k) ans = max(g,ans);
    }

    cout << ans << endl;
    return 0;
}