ひらめの日常

プログラミングと読書と

Codeforces Round #582 (Div. 3) D. Equalizing by Division

問題はこちら

codeforces.com

長さ  n の 配列  a が与えられる。一回の操作によって任意の要素一つ、  a_i を2で割って切り捨てることを行う。 (つまり、  a _ {i} :=\left\lfloor\frac{a _ i}{2}\right\rfloor

 k 個の等しい要素を  a 中で得るには、何回操作を行えばいいか、その最小値を求めよ。

考え方

全ての  a_i を、0になるまで割る2して、各値を得るまでにかかった作業回数を記録しておく。コードの方がわかりやすいと思うのでその部分を抜粋する。

while (a[i] > 0) {
    mp[a[i]].emplace_back(t);
    a[i] /= 2;
    t++;
}

そして、存在している全ての要素に対して以下を行う。

  • 素数 k よりも以上の場合は、要素をsortする。
  • その後、  k 番目までの必要回数の和を計算する。

以上を行なって、 k 個以上作ることのできる要素の最小操作回数を求め、それが答えとなる。

なぜこれで間に合うのか?

最初は、全ての要素に対して、各要素を得るための操作回数をsortするために間に合わないと思っていたがそんなことはなかった。

以下は簡略化のために  O(max(a)) = O(n) として話を進める。 まず、全ての要素数 O(nlog(n)) である。(一つの要素は2で割っていくので  log(n) 個存在し、それが  n 個あるため。)

次に、その各値に対応する操作回数をsortする部分だが、sortする対象の合計個数も高々  O(nlog(n)) 個である。よって、各値に対応する操作回数を全てsortしたとしても計算量は  O(xlog(x)), x = nlog(n) で抑えられ、時間内に間に合う。

基本的に、合計個数  n あるものを、それぞれ  A, B, C... 個を要素とするグループに分けたときを考える。計算量 O(n) 以上となる同じ操作hogeを各グループに対して行なった合計の計算量は、合計個数全てに対する計算量で抑えられる。つまり、  O(hoge(A)) + O(hoge(B)) + O(hoge(C)) + ...\ =\ O(hoge(n))

hogeの計算量が  O(log(n)) とかだと、グループに分けない方が計算量が小さいはず。今回は hoge に相当するのが sort なので、 O(n) 以上の計算量である。)

解答

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
#define rep(i, n) for(ll i = 0; i < (ll)(n); i++)
#define each(i, mp) for(auto& i:mp)

int main() {

    ll n, k;
    cin >> n >> k;
    vector<ll> a(n);
    rep(i, n) cin >> a[i];

    unordered_map<ll, vector<ll>> mp;
    rep(i, n) {
        ll t = 0;
        while (a[i] > 0) {
            mp[a[i]].emplace_back(t);
            a[i] /= 2;
            t++;
        }
    }

    ll ans = 1e18;
    each(e, mp) {
        vector<ll> times = e.second;
        if (times.size() >= k) {
            sort(times.begin(), times.end());
            ll now = accumulate(times.begin(), times.begin() + k, 0LL);
            ans = min(ans, now);
        }
    }
    cout << ans << '\n';
    return 0;
}

Submission #59801874 - Codeforces