ABC141D - Powerful Discount Ticketsを解いてみた

40代文系社内SEがAtCoder Beginner Contest 141 D - Powerful Discount Ticketsを解いてみました。

まず問題文をみます。

半額　小数点以下切り捨て　になる割引券という夢のような設定です。それって最終的には0円もありえる？と思ったら入力例3にありました。

割引前の金額に対して割合で値引きをするので、割引前の金額が高いものからチケットを使用していきたいところです。問題には、「順番に購入する」とありますが、購入前のシミュレーションと考えれば購入順は気にしなくてよさそうです。

そのような要件に応えるようなアルゴリズムとしては、優先度付きキューがあります。優先度つきキューなら最大値の取り出しと値の追加が$O(\log N)$で完了します。

次に制約をみます。

品物の個数$N$と割引券の枚数$M$がともに最大$10^5$です。

割引券ごとに優先度つきキューから商品を取り出して割引を適用してキューに戻すという操作を繰り返しても、$O(M \log N)$で完了します。

以上から、以下の手順で実装します。

1. 最大値を返す優先度付きキューを用意する
2. 優先度付きキューに商品を追加する
3. 割引券ごとに、優先度付きキューから商品を取り出して割引を適用してキューに戻す操作を繰り返す
4. 優先度付きキューの値の合計を出力する

(defun solve ()
  (let* ((n (read))
         (m (read))
         (heapq (make-array 0 :fill-pointer 0 :adjustable t)))
    (loop repeat n
          for ai = (read)
          do (heap-push heapq ai))
    (loop repeat m
          for ai = (heap-pop heapq)
          do (heap-push heapq (floor ai 2)))
    (format t "~A~%" (reduce #'+ heapq))))


(defun shift-down (q idx)
  (let ((idx-l (+ (* idx 2) 1))
        (idx-r (+ (* idx 2) 2))
        (idx-next idx))
    (when (and (< idx-r (length q))
               (> (aref q idx-r) (aref q idx-next)))
      (setf idx-next idx-r))
    (when (and (< idx-l (length q))
               (> (aref q idx-l) (aref q idx-next)))
      (setf idx-next idx-l))
    (when (/= idx-next idx)
      (rotatef (aref q idx) (aref q idx-next))
      (shift-down q idx-next))))

(defun heap-push (q item)
  (vector-push-extend item q)
  (loop with idx = (1- (length q))
        for parent-idx = (floor (1- idx) 2)
        while (and (> idx 0)
                   (> (aref q idx) (aref q parent-idx)))
        do (rotatef (aref q idx) (aref q parent-idx))
           (setf idx (floor (1- idx) 2))))

(defun heap-pop (q)
  (cond ((zerop (length q))
         nil)
        ((= (length q) 1)
         (vector-pop q))
        (t
         (let ((ans (aref q 0)))
           (setf (aref q 0) (vector-pop q))
           (shift-down q 0)
           ans))))

(defun heap-top (q)
  (if (zerop (length q))
      nil
      (aref q 0)))

(solve)

Common Lispには優先度付きキューがないので、自前で実装しました。