🙆‍♀️

ABC 321 D - Set Menuのメモ

2023/09/26に公開

コンテストでは解けなかったので、後日解いたものです。
https://atcoder.jp/contests/abc321/tasks/abc321_d

あまり時間がないので手短に…

考察

この問題はmin(s,P)の部分によって難しくなっています。これがなければ、たとえばBの合計値を求めて、それにAi × Mを足していくだけで求められるのですが、それができなくなっています。

足してもPを超えない部分についてはこのような考え方でいいのですが。そこで、足してもPを超えない部分と、超える部分の2つに分けて考えます。
超えない部分と超える部分の境目がわかれば、(超えない部分の合計値) + Ai × (超えない部分の長さ)を求めて、(足すとP以上になる部分の長さ) × Pを足すだけで求められます。

なのでその境目を見つける問題と整理できます。それを探すためには二分探索が使えます。線形探索でも見つけることはできますが、遅すぎて間に合わないので二分探索を使います。
Bをソートしておく必要がありますが、Bについて「PからAiを引いた数」を二分探索で探します。「PからAiを引いた数」とは、Aiを足すとP以上になる最初の数で、つまりこれがまさに見つけたい境目です。
ただ、Bの中にちょうどその値があるとは限りません。見つからない場合は、その値より大きい最初の数の位置を見つける必要があります。これは、言語によっては組み込みの二分探索メソッドでも実現可能です。
C++ならlower_bound()が使えますが、私が使っているKoltinだと普通に二分探索のために使うbinarySearch()メソッドでそれを見つけることができます。
今回、人の提出コードを見ていて初めて知ったのですが、見つからなかった場合の戻り値(負の数)に1を足してから符号を反転すると、仮にその値があるのであれば返っていたはずの値がわかります。今回の問題でまさに欲しい値です。

https://note.com/kirimin_chan/n/n3b5c9a0d4290

その値さえわかれば、あとは上で書いた(超えない部分の合計値) + Ai × (超えない部分の長さ)を求めて、(足すとP以上になる部分の長さ) × Pを足す計算をすればOKです。

ただ、(超えない部分の合計値)をいちいち求めていたら遅くて間に合わないので、累積和を使います。

実装

fun main() {
    val (n, m, p) = readln().split(" ").map { it.toLong() }
    val a = readln().split(" ").map { it.toLong() }
    val b = readln().split(" ").map { it.toLong() }
        .sorted()

    val s = LongArray(m.toInt() + 1)
    for(i in 0 until m.toInt()) {
        s[i+1] = s[i] + b[i]
    }

    var ans = 0L
    for(i in 0 until n.toInt()) {
        val diff = p - a[i]

        val tmp = b.binarySearch(diff)
        val index = if(tmp < 0) -(tmp + 1) else tmp

        ans += (m - index) * p
        ans += a[i] * index + s[index] - s[0]
    }

    println(ans)
}

このように二分探索で位置を探し、見つかった位置情報を使って累積和で高速に区間の合計値を求めるというのは頻出のパターンらしいですね…
覚えておきたいと思います。

一応、以前もその組み合わせは見たことがあったのですが、その問題は今回の問題よりはだいぶ高度で、あまり理解できませんでした。
https://zenn.dev/dhirabayashi/articles/557bbf51c35edf#d---iroha-and-haiku-(new-abc-edition)

今回は、二分探索と累積和の組み合わせという意味だとかなりシンプルな問題なのでわかりやすかったです。
(執筆時間: 20分23秒)

Discussion