【Python】累積和について勉強したことまとめ

スポンサーリンク
AtCoder
スポンサーリンク

最近、累積和について勉強したのでメモとして残しておきます。

今回は、pythonで累積和を使った問題(AtCoder Beginner Contest154 問題D)を解く方法をまとめました。

スポンサーリンク

累積和とは?

累積和を使うことで、配列のある区間の総和をすばやく求めることができます。

例えば、配列Aの紫色の部分の総和を求める場合を考えます。

通常では、この部分の総和を求める時、

4 + 5 + 6 + 4 = 19

という計算で求めると思います。

続いて、累積和を使った方法で総和を求めて見ましょう。

まず、配列のi番目までの総和を、別の配列(配列B)に保存しておきます。

次に、配列Bを使って、紫色の区間の総和を求めます。

求めたい区間の右端の値(配列B)から、左端の値(配列B)を引くことで、紫の部分の総和を計算することができます。

配列Bの6番目 – 配列Bの2番目 = 24 – 5 = 19

紫の区間の数字を一つずつ足した時の値と同じ値が得られました。

累積和に関して詳しくはこちらのサイトがとても参考になります。

累積和を何も考えずに書けるようにする!

問題

実際に累積和を使った問題を解いてみます。

N 個のサイコロが左から右に一列に並べてあります。左から i 番目のサイコロは 1 から pi までの pi 種類の目がそれぞれ等確率で出ます。

隣接する K 個のサイコロを選んでそれぞれ独立に振ったとき、出る目の合計の期待値の最大値を求めてください。

制約

  • 1≤K≤N≤200000
  • 1≤pi≤1000
  • 入力で与えられる値は全て整数

例として、N=5, K=3, サイコロの出る目の種類はそれぞれ、[1, 2, 2, 4, 5]の場合を考えます。

期待値を求める

まず、それぞれのサイコロの期待値を求めます。

P種類の目が等確率で出るサイコロの出る目の期待値は、(1+P)/2です。

前処理

この問題も、区間ごとの総和を求めたいので、i番目(0 <= i >= N)までの総和(累積和)をあらかじめ計算しておきます。

pythonでは、itertoolsというライブラリのaccumulateを使うことで簡単に累積和を計算できます。

コード全体はこちらです。

計算量に関して

もし、累積和を使わずにこの問題を解くとします。

このように、区間をずらしながら、それぞれの区間ごとに総和を計算し、その中の最大値を求めることになると思います。

そうすると、

  • 区間ごとの総和を計算 O(N)
  • 区間をずらす O(N)

という二つの処理を同時に行うので、O(N2)の計算量になってしまいます。

一方、累積和を使った方法では、

  • 累積和を計算(前処理)O(N)
  • 区間をずらす O(N)

という処理に分けているので、計算量がO(N)で済みます。

タイトルとURLをコピーしました