【競技プログラミング】Rubyで累積和を計算する方法

累積和を計算するコードをすらすらかけますか？競技プログラミングの序盤は累積和を計算する問題が頻繁に出題されます。私も累積和が必要となる場面に何度か遭遇したので、勉強した内容をアウトプットします。

今回は一番簡素な１次元の配列の場合をRubyで実装する方法と考え方を解説します。

累積和とは

累積和とは配列において、i番目の際に最初からi番目まで計算したものです。

例えば、

1番目は、[1番目]を計算したもの、

2番目は、[1番目]＋[2番目]を計算したもの、

3番目は、[1番目]+[2番目]+[3番目]を計算したもの、

i番目は、[1番目]+[....]+[....]+[i番目]まで計算したもの

です。

下の図にある配列wの累積和を計算する場合、累積和sのs[3]は、4+5+2をしたものです。

s[3]は、別の表し方をすると、s[3] = s[2] + w[2]となります。

一般式は、s[i] = s[i-1] + w[i-1]です。

別の表現ですと、

s[i+1] = s[i] + w[i]でもOKです。

累積和の計算量

配列の長さN回分、ループするので、O(N)です。

累積和を使う場面

配列のある区間の和を求める時に便利です。

例えば、下の赤い枠の和と青い枠の和のどっちが大きいかを計算する場合です。

素朴に計算すると、w[1] + w[2] + w[3] = 8 と、w[5] + w[6] + w[7] = 9を比べることになります。今回は、枠の幅3、枠の種類2種類でしたが、計算量は枠の幅×枠の種類になります。O(NM)。枠の幅がめっちゃ広い場合はそれに応じて計算量も多くなります。

累積和を用いると、s[4] – s[1] = 8と、s[8] – s[5] = 9を比べます。今回は、枠の幅はどう広がろうと、先端 – 末端をするので、2×枠の種類となります。O(M)

Rubyで累積和を使用してみる場合

全体像はこのようになる。

w = [4,5,2,1,6,4,2,3,2]
section = [
  [1,3],
  [5,7]
]
s = Array.new(w.length+1)

s[0] = 0

0.upto(w.length-1) do |i|
  s[i+1] = s[i] + w[i]
end

# もしくは
# 1.upto(w.length) do |i|
#   s[i] = s[i-1] + w[i-1]
# end

max = 0
section.each do |first, second|
  if s[second+1] - s[first] > max
    max = s[second+1] - s[first]
  end
end

puts max
# max => 9

まず、累積和を求める部分、s[0]は0で初期化しておいて、s[i+1] = s[i] + w[i]を計算する。

s = Array.new(w.length+1)

s[0] = 0

0.upto(w.length-1) do |i|
  s[i+1] = s[i] + w[i]
end

累積和を使って、区間の和を求める部分

max = 0
section.each do |first, second|
  if s[second+1] - s[first] > max
    max = s[second+1] - s[first]
  end
end

２次元配列[[1,2],[3,4]]をeachで回して、|A,B|のようにすると、1週目のAには1、Bには2が入ります。

maxより大きかったらmaxを置き換えるようにすれば、最大値を求めることができそうです。

元の配列を更新していく場合

先ほどとは打って変わって、配列sを生成しない方法を考えます。

赤枠(18)を求める場合は、元々の配列の値(6)と、一個前の累積和(12)を足せば良さそうです。

一般式は、w[i+1] = w[i+1] + w[i]です。

もしくは、w[i+1] += w[i]となります。

これも同じく累積和となります。

【不安解消】Rubyでも競技プログラミングで問題が解ける→茶色に行けました!

累積和はこのサイトでも分かりやすく解説されています。