🌊

ABC246 C - Couponのメモ

2023/02/21に公開

目についたC問題を解いてみるの回です。解けたけど異様に苦労しました。
https://atcoder.jp/contests/abc246/tasks/abc246_c

考察

支払う金額を最小にするには、可能な限り高い商品に対してクーポンを使う必要があります。なので基本的には商品の値段で降順ソートしてクーポンを適用していくことを考えます。ただ、商品が0円になるまでクーポンを使ってしまうと最後の一枚が無駄になるかもしれないので、端数に対してはクーポンを使わないようにします。
とりあえず端数が出るまでは使って、端数は商品一覧に戻して再度降順ソートして次の高い商品に対してクーポンを使う、を繰り返せばいいですね。(元の金額はもちろん一覧から削除する)

もちろんそのためにはそれなりに賢いデータ構造を使う必要があります。取り出し、削除、追加およびその後のソートが遅かったら間に合いません。
これには、優先度付きキューを使うとよさそうだと思いました。追加時に自動的にソートされます。先頭要素に対する取得と削除、ソートも含めた追加の操作がO(logN)で実行できます。
JavaではPriorityQueueとして実装されているのでこれが使えます。普通は昇順のソートされますが、コンストラクタにComparatorを渡せばその通りにソートされるので、降順ソートするようなComparatorを渡せばいいです。

なお、この問題の制約としては商品の数に対して線形時間は許容されますが、クーポンの数に対して線形時間は許容されません。
(何度かTLEを出してしかもしばらく理由がわからなかったので、我ながら…うん)

実装

最初にACしたコードがわかりづらかったので、コメントを追加して変数名を見直しました。
あとなぜかクーポンをチケットと書いてしまった箇所をこっそり直した

import java.util.*
import kotlin.math.max

fun main() {
    val (n, k, x) = readLine()!!.split(" ").map { it.toLong() }
    val a = readLine()!!.split(" ").map { it.toLong() }

    // Comparator.reverseOrder()は現時点のAtCoderのジャッジサーバだとコンパイルエラーになる(JVM target 1.6で実行されるため)
    val productPriceQueue = PriorityQueue<Long> { i1, i2 ->
        val diff = i2 - i1
        if(diff < 0) {
            -1
        } else if(diff > 0) {
            1
        } else {
            0
        }
    }

    productPriceQueue.addAll(a)

    var currentCouponCount = k
    for(i in productPriceQueue.indices) {
        // クーポンがなくなっていたら抜ける
        if(currentCouponCount == 0L) {
            break
        }

        // 現状最も高い商品の値段を取り出す
        val maxPrice = productPriceQueue.poll()

        // 消費するクーポンの枚数。切り捨て除算により、端数が余る枚数を得られる
        var useCouponCount = maxPrice / x

        // クーポン枚数が足りなければ、ある分だけを使う
        if(useCouponCount > currentCouponCount) {
            useCouponCount = currentCouponCount
            currentCouponCount = 0
        } else {
            currentCouponCount -= useCouponCount
        }

        // クーポンを可能な限り使った後の金額をキューに戻す(ここで負の数になることはない)
        productPriceQueue.add(maxPrice - useCouponCount * x)
    }

    // 端数が余っているかもしれないのでもう一周する
    for(i in productPriceQueue.indices) {
        // クーポンがなくなっていたら抜ける
        if(currentCouponCount == 0L) {
            break
        }

        // 現状最も高い商品の値段を取り出す
        val maxPrice = productPriceQueue.poll()

        // この時点ではクーポンが尽きてない限りは端数しか残っていないので、1枚ずつ使う
        // 端数に対してクーポンを使うため、負の数になることを避けるための考慮が必要
        productPriceQueue.add(max(maxPrice - x, 0))
        currentCouponCount--
    }

    // 残った金額の合計値が答え
    println(productPriceQueue.sum())
}

ポイントがいくつか。まず降順ソートするために本来はComparator.reverseOrder()を使えるのですが、ジャッジサーバ起因により使えないのでラムダ式で同等のものを作っています。

アルゴリズムとしては上で書いた通りですが、商品を一通りなめてもまだ端数が残っているかもしれないのでもう一周するようにしています。微妙といえば微妙ですが…
クーポンの一覧を走査したくなりますが、制約的に許されないのは前述の通り。(まあ、やってしまいましたが…)

二周目については、端数しか残っていない(最大でも1枚しかクーポンを使えない)かもしくはクーポンが残っていないかなので、商品一覧を走査すれば問題ないです。

この実装だとキューに0を書き戻す可能性がありますが…実行結果には影響ないです。0円に対してクーポンを適用しようとするかもしれませんが、負の数にならないようにしていますし、最大値を取り出しているのに0が取れるということは答えがもう確実に0なので、クーポンを無駄にしても結果は変わりません。

あとは、入力値はIntの範囲内ですが計算の途中経過でIntだとオーバーフローするかもしれないのでLongで扱います。(ミスりました)
また、クーポンの残数以上にクーポンを使おうとしてはいけないのでその考慮も必要です。(これもミスった)

感想

実装をミスり過ぎですね。考察自体はそんなに苦労しなかったのですが。実装力まで課題になるとは…厳しい。

(執筆時間: 50分38秒)

Discussion