2019-05-15

AtCoder ABC 116 D - Various Sushi (青色, 400 点)

AtCoder AtCoder400点 ABC-D N個のペア値の問題最適化テク:最適解の形を考える最適化テク:緩和しても良い最適化テク:解を変形していく(最適性を失わずに) データ構造 setの上手な使い方差分更新 N個からK個を選ぶ設定の問題 Greedy:ある量を決めると残りが決まっていくある値を固定して考える Greedy priority_queue 探索順序を工夫して解く番兵法最大スコア青色diff 非自明な比較関数でソートする最適化問題

学び多き問題。
僕にとっては後半のデータ構造パートが苦戦を強いられ、本当に勉強になった！

問題へのリンク

問題概要

$N$ 個の寿司があって、それぞれネタ $t_i$ と美味しさ $d_i$ をもっている。この中から $K$ 個の寿司を選びたい。選んだ寿司集合のスコアは

選んだ寿司の美味しさの総和
選んだ寿司に含まれるネタの種類数を $x$ として $x \times x$

の合計値で決まる。スコアの最大値を求めよ。

制約

$1 \le K \le N \le 10^{5}$
$1 \le t_i \le N$

考えたこと

この手の最適化問題では

ある量を決めるとどうすべきかが Greedy に決まらないか
「こういうものだけ探索すればよい」というのが絞れないか

をひたすら考えることになる。こういう考察を意識的にやることは高難易度でも有効なイメージ。この問題では

選ぶネタの種類数 $x$

を固定したくなる。 $x$ を決めると最適解が自然に決まるのではないかと考えたくなる。ここでしばしばやる注意点として

選ぶネタの種類がちょうど $x$ でなければならない

としてしまうと考察しづらくなることが多い。より柔軟に

「選ぶネタの種類が $x$ 種類以上になる場合についての、美味しさ」の総和の最大値を求め、 $x \times x$ を足したもの」を各 $x$ について求めて、その最大値を求める

とした方がやりやすい。このとき、 $x$ 種類の場合についての最適解が実際には $y (\gt x)$ 種類だった場合は、種類数ボーナススコアが損をするわけだが、気にしなくて良い。なぜなら、どちらにしても $y$ 種類の場合について探索したときに、その解を検討することになるからだ。

こういう風に「ちょうど」を「以上」にした方がいい / しても大丈夫、という発想はメチャクチャよく見るパターンではある。

解の形

以上から我々の解きたい問題は $x$ 種類以上選ぶときの美味しさの総和を最大化することに帰着された。

そうと決まれば、 $N$ 種類ある中からどの $x$ 種類を選ぶかを考えることになる。このままでは ${}_{N}{\rm C}_{x}$ 通りあって手に負えない。なにかしら Greedy に決まる要素があるはずである。

まず、各ネタについて美味しさを大きい順に並べておいて、さらに各ネタを「美味しさの最大値」が大きい順に並べてみる (下図はネタの種類の番号も「美味しさの最大値」が大きい順に振り直したことを想定)

ここでふと思い浮かぶのは

実は美味しさの最大値が大きい順に $x$ 種類選ぶ

で良いのではということである。少し考えるとこれが実際に正しいことがわかる。というのも、例えば下図のように、もし「美味しさの最大値」を大きい順に $x$ 種類選ぶのではなく、どこか飛ばしてしまったとしたら、それを飛ばさないように選ぶように変更しても解が悪化することがない (同点でない限りは良くなる) からだ。ここで、下図では結果的に種類が増えて $x+1$ 種類になっているが、それも気にしなくて OK。

最適化問題では、こういう「解を悪化させずにこういう風に変形できるから、こういう解だけ調べればよい」という考察はものすごくよくある！！！！！

よって、

$N$ 種類の中から、「美味しさの最大値」が大きい順に $x$ 個選び、その $x$ 種類の中で美味しさが最大のものをまずは確保する
残りの $K - x$ 個については、残った中から大きい順に $K - x$ 個とれば OK (このとき、選んだ $x$ 種類以外からも選んでしまってもよい)

ということがわかった。

計算量落とし

この時点で、ひとまず $O(KN\log{N})$ にはなっているが、高速化が必要だ。ここは、こどふぉで鍛えた人たちは強い。僕は超苦戦した。つまり $x$ 種類選んだ状態と $x+1$ 種類選んだ状態との間の差分をいかに素早く更新できるようなデータの持ち方を工夫するか、ということになる。

色々考えてみると、 $x = K$ から順番に $x = K, K-1, K-2, \dots$ という順序でやった方がよさそう。

$x = K$ のときは、上位 $K$ ネタについての「美味しさの最大値」を採用しておく
のときの解からのときの解を構成するためには、番目のネタの「美味しさの最大値」と、「現在の個には選ばれていないお寿司の美味しさの最大値」とを比較して、
- $v \ge w$ のときは実は何もしなくて OK
- $v \lt w$ のときは $v$ を破棄して $w$ を入れる

という風にすれば良くて、こういうことができるデータ構造として priority_queue がある。すなわち

que：選ばれていないお寿司の美味しさ (最大値を取得できるようにしておく)

を用意しておいて、常に v と w = que.top() とを比較して、交換が必要になったら

que.pop()
que.push(v)

をすれば OK。なお人によって本当に色々なデータ構造を使った模様。set を使ってもできる。いずれにせよ、これで $x+1$ から $x$ への変更を $O(\log{N})$ で実現できることとなった。総じて計算量は $O(N\log{N})$ になる。

#include <iostream>
#include <queue>
#include <algorithm>
#include <numeric>
using namespace std;
const long long INF = 1LL<<40;

int main() {
    long long N, K; cin >> N >> K;
    vector<vector<long long> > s(N); // s[i] := 種類 i の美味しさの集合
    for (int i = 0; i < N; ++i) {
        long long t, d; cin >> t >> d; --t;
        s[t].push_back(d);
    }

    // 各種類について、美味しい順に並べる
    for (auto &v : s) {
        sort(v.begin(), v.end(), greater<long long>());
        if (v.empty()) v.push_back(-INF); // 番兵
    }

    // 「美味しさの最大値」の大きいネタ順
    sort(s.begin(), s.end(), [&](vector<long long> a, vector<long long> b) {
            return a[0] > b[0];});
 
    // まず K 種類選ぶ場合
    long long cur = 0;
    priority_queue<long long> que;
    for (int i = 0; i < K; ++i) {
        cur += s[i][0];
        for (int j = 1; j < s[i].size(); ++j) que.push(s[i][j]);
    }
    for (int i = K; i < N; ++i) for (int j = 0; j < s[i].size(); ++j) que.push(s[i][j]);

    // 走査
    long long res = cur + K*K;
    for (long long x = K-1; x >= 1; --x) {
        long long v = s[x][0];
        long long w = que.top();
        if (v < w) {
            que.pop(); cur += w;
            que.push(v); cur -= v;
        }
        res = max(res, cur + x*x);
    }
    cout << res << endl;
}