Rubyで累積和を計算する方法

  • このエントリーをはてなブックマークに追加

競技プログラミングで累積和を計算する場面に何度か遭遇するので、勉強した内容を解説してアウトプットします。

今回は一番簡素な1次元の配列の場合をやります。

累積和とは

配列において、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

2次元配列[[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]となります。

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

  • このエントリーをはてなブックマークに追加

SNSでもご購読できます。

スポンサードリンク

コメント

コメントを残す

*