LoginSignup
1
1

Pythonで累積和〜愚直解から最適解まで〜

Last updated at Posted at 2023-10-28

下記の例題は累積和の基礎となるAtcoderの問題です。
今回はこの例題をもとに、累積和を求める愚直解から最適解までのステップを確認していきます。

愚直解

各質問に対して、毎回LからRの来場者数の和を直接計算する

この方法ではQ回の質問に対して、毎回L日からR日までのの和を求める計算をします。
その場合、計算量は最大O(N * Q)になります。
N, Qともに取りうる最大値は10 ^ 5なので実行時間制限2秒を超えてしまいます。

累積和を事前に求める(直接計算)

1つ目の方法は同じ日数間の合計来場者数を求める場合でも、再度計算する必要がありました。
そのため計算量が大きくなることが欠点でした。

今回は事前にそれぞれの日数の合計来場者数を求めることとします。その場合、同じ日数間の合計来場者数はO(N)からO(1)に削減できます。

例えば、A = [11, 46, 47, 77, 80]のリストがある場合、この累積和を計算してみます。

  • S1 = 11
  • S2 = 11 + 46 = 57
  • S3 = 11 + 46 + 47 = 104
  • S4 = 11 + 46 + 47 + 77 = 181
  • S5 = 11 + 46 + 47 + 77 + 80 = 261

この配列の大きさであれば手計算でも可能ですが、配列の大きさがNの場合、S1を求めるのに1個の数の足し算、S2を求めるのに2個の数の足し算、SNを求めるのにN個の数の足し算をする必要があります。合計計算回数は N * (N - 1) / 2回になり、Nが大きい場合に実行時間に間に合いません。
そのため次の最適解である累積和の計算方法を取ることにします。

最適解

累積和を事前に求める(一つ前に求めた累積和を使って次の累積和を求める)

以下の方法を取ることで、計算量O(N)で累積和を求めることが可能になります。

  • S0 = 0とする
  • i = 1, 2, 3, …., Nとした場合、Si = Si-1 + Aiを計算して累積和Siを求める

A = [11, 46, 47, 77, 80]のリストがある場合、この累積和を計算してみます。

  • S1 = S0 + A1 = 11
  • S2 = S1 + A2 = 57
  • S3 = S2 + A3 = 104
  • S4 = S3 + A4 = 181
  • S5 = S4 + A5 = 261

例題では累積和を使ってQ回の質問に答えるので計算量はO(Q + N)で計算できます。
ここで例題の注意点として、L日からR日間の累計来場者数を計算する場合、
(Rの累計来場者数)- (L-1日目の累計来場者数)を計算する必要があります。
この問題ではL日目の来場者数も含める必要があるため、L-1としています。

最適解による回答

N, Q = map(int, input().split())
A = list(map(int, input().split()))
L = [None] * Q
R = [None] * Q
for q in range(Q):
    L[q], R[q] = map(int, input().split())

# 累積和の計算
S = [None] * (N + 1)
S[0] = 0
for i in range(N):
    S[i + 1] = S[i] + A[i]

for q in range(Q):
    print(S[R[q]] - S[L[q] - 1])

Pythonの場合、一次元の累積和は**itertools.accumulate**のライブラリを使うことで以下のように簡単に計算することができます。

from itertools import accumulate

N, Q = map(int, input().split())
A = list(map(int, input().split()))
L = [None] * Q
R = [None] * Q
for q in range(Q):
    L[q], R[q] = map(int, input().split())

# 累積和の計算
S = list(accumulate([0] + A))

for q in range(Q):
    print(S[R[q]] - S[L[q] - 1])
1
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
1
1