ABC116 D - Various Sushi

こういう、適切な貪欲をやる問題が苦手すぎる…。
D - Various Sushi

問題

N個の寿司があり、i個目の寿司の種類はt_i、美味しさはd_iである。この中からK個選び、美味しさd_iの総和とt_iの種類数の2乗の和を最大化したい。最大値を求めよ。

解法

寿司を美味しさの順に並べたとき、各種類について一番美味しさが高いものをinitiator、それ以外のものをfollowerと呼ぶことにする。
最適な選び方を一つ固定する。これはどういう形をしているだろうか? まず、各種類について、その種類の寿司を取ったのであれば必ずinitiatorは取っている。また、選んだfollowerよりも価値の高い選んでいないinitiatorがある場合、followerをinitiatorと取り替えることによりスコアが高くなるので、そのようなことはないとして良い。
よって、最適戦略は上からinitiatorを何個か取って、そのあとにfollowerを何個か取った形をしている。上から取るinitiatorの個数を全探索すればよい。
initatorを新しく取るたびに一番美味しさの低いfollowerを捨てる操作が高速にできればよい。これは優先度キュー、vectorなどでできる。

登場する典型

  • 貪欲
  • 最適解の形を考える

実装上の注意点

  • 新しいinitiatorを追加するたびに、2 * (今まで取ったinitiatorの個数) + 1だけスコアに加算するのを忘れない!

提出: Submission #4059339 - AtCoder Beginner Contest 116 (Rust)

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        k: usize,
        td: [(usize1, i64); n],
    }
    let mut dt = vec![];
    for (t, d) in td {
        dt.push((d, t));
    }
    dt.sort(); dt.reverse();
    let mut que = BinaryHeap::new();
    let mut seen = vec![false; n];
    let mut tot = 0;
    let mut kind = 0;
    for i in 0 .. k {
        let (d, t) = dt[i];
        tot += d;
        if seen[t] {
            que.push((-d, t));
        } else {
            seen[t] = true;
            tot += 2 * kind + 1;
            kind += 1;
        }
    }
    let mut ma = tot;
    for i in k .. n {
        let (d, t) = dt[i];
        // No point of picking it. Skip.
        if seen[t] { continue; }
        let profit = d + 2 * kind + 1;
        if let Some((tap, _ris)) = que.pop() {
            let tap = -tap;
            tot += profit - tap;
            seen[t] = true;
            kind += 1;
        }
        ma = max(ma, tot);
    }
    puts!("{}\n", ma);
}

まとめ

最近のABCは結構定跡の把握漏れに気付かせてくれる気がする。