https://paiza.jp/poh/enkoi の問題３

数列 $a_0, a_1,\dots,a_n$が与えられた時に、連続する$t$項の和$a_i,a_{i+1},\dots,a_{i+t-1}$の最大値を求める問題を$O(n)$で解けという話。

手続き的に書くのは簡単で、配列のインデックスを進めるたびに当該の区間から出て行った項を減算して新しく入った項を加算してやれば各区間の和が順次的に定数時間で求められるから、区間の和の配列($s_0,s_1,\dots$)とその最大要素が$O(n)$で求められる。区間和$$s_i=a_i+a_{i+1},\dots,a_{i+t-1}$$ は漸化式 $$s_{i+1}=s_i-a_i+a_{i+t}$$を満たすから
数列 $a_0,\dots,a_{11}$ = {0 123 222 21 1 332 22 99 3 444 24 10}に対してt=3の場合で、
(0+123+222=345), (345-0+21=366), (366-123+1=244),... とできるわけ。問題はインデックス操作をせずにこれを書くにはどうしたらいいか、である。実はわりと簡単で：
$$ s_i = (a_{i+t-1}-a_{i-1})+s_{i-1} = (a_{i+t-1}-a_{i-1})+(a_{i+t-2}-a_{i-2}) + s_{i-2} = \dots$$ なので $$s_{i+1} = \sum_{k=0}^{i} (a_{k+t} - a_{k}) + \sum_{k=0}^{t-1} a_k$$ になる。

たとえば次の数列：

1, 3, 6, 8, 2, 7, 7, 9, 0, 4, 2, 7, 1,...

に対して3項連続の区間の和を計算したければ、最初の3項を取り除いた数列を用意し、元の数列をそこから各項ごとに減算する：

8, 2, 7, 7, 9, 0, 4, 2, 7, 1,...
1, 3, 6, 8, 2, 7, 7, 9, 0, 4, 2, 7, 1,...
-----------------------------------------
7,-1, 1,-1, 7,-7,-3,-7, 7, -3,...

この数列の累積和を取り（最初に0を付け加えておく）：

0, 7, 6, 7, 6, 13, 6, 3, -4, 3, 0,...

最後に各項に、最初に取り除いた3項の和 1+3+6=10を足すと

10, 17, 16, 17, 16, 23, 16, 13, 6, 13, 10,...

区間和の数列ができる。あとは最大値を探すだけ。これらの操作はすべて$O(n)$なので、全体も$O(n)$で終わる。

fmax t as =  maximum $ scanl (+) (sum$take t as) (zipWith (-) (drop t as) as)

というわけで恒例のウンコード

こんなウンコードでも大丈夫なのか……

import Control.Applicative
import Data.List

main = do
  (h:ts) <- filter (not.null) <$> lines <$> getContents
  let [t,n] = map (read::String->Int) $ words h
  let nums = map (read::String->Int) ts
  print $ fmax t nums
    where 
      fmax t as =  maximum $ scanl (+) (sum $ take t as) (zipWith (-) (drop t as) as)