ABC154D - Dice in Lineを解いてみた

40代文系社内SEがAtCoder Beginner Contest 154 D - Dice in Lineを解いてみました。

まず問題文を見ます。

期待値を求め、合計を計算し、さらに最大値を求めると。$N$個から$K$個の合計を計算することになるので、愚直に計算すると$O(NK)$となるのでおそらくTLEになるでしょう。

次に制約を見ます。

$1 \leq N \leq K \leq 200000$ということで、やはり期待値の合計を計算するところは工夫が必要です。正$1000$面体のサイコロはほぼ球になる気がします。出目がわかりにくそう。

さて、解き方についてですが、まず期待値についてはサイコロの目をすべて合計して面数で割ります。これも愚直に計算すると$O(p)$になり、すべてのサイコロの期待値を計算すると$O(Np)$となり間に合いません。

$1$から$K$までの合計を計算するには、$1 + K \times K \div 2$を使えば$O(1)$で求められます。例えば、一般的な正六面体の場合は、$(1 + 6) \times 6 \div 2 = 21$となります。

次に$N$個から$K$個の合計を計算するところですが、「連続する$K$個」というところがポイントです。まず先頭$K$個の合計を求め、合計1とします。合計2を求めるには、合計1に$K+1$個目の値を足し、1個目の値を引くことで$O(1)$で求められます。

例えば、入力例1では以下のようになります。

合計1: $1.0 + 1.5 + 1.5 = 4.0$
合計2: $4.0 + 2.5 - 1.0 = 5.5$
合計3: $5.5 + 3.0 - 1.5 = 7.0$

以上より、以下の手順で解きます。

1. 各サイコロの期待値を計算する
2. $K$個の期待値の合計を計算する
3. 期待値の合計の最大値をとる

(defun solve ()
  (let* ((n (read))
         (k (read))
         (p (make-array n :initial-contents (loop repeat n collect (read))))
         (q (make-array n)))
    (loop for i from 0 below n
          for p-i = (+ (aref p i) 0d0)
          do (setf (aref q i) (if (= p-i 1)
                                  1
                                  (/ (* (1+ p-i) p-i) 2 p-i))))
    (loop with a = 0
          for i from 0 below n
          for qi = (aref q i)
          maximize (progn
                     (incf a qi)
                     (when (>= (- i k) 0)
                       (decf a (aref q (- i k))))
                     a)
            into it
          finally
             (format t "~,10F~%" it))))

(solve)

0d0は、倍精度浮動小数点数で計算するために使用しています。