🍡

[鉄則A06] How Many Guests?

2024/04/19に公開

https://atcoder.jp/contests/tessoku-book/tasks/math_and_algorithm_ai

累積和の基本が学べる問題。

問題を要約する

  • 日毎のイベント来場人数の記録がある
  • 期間 L..R の合計来場人数は?

シミュレーション

来場人数の、

A = [2, 3, 4, 5]

から累積和の、

S = [0, 2, 5, 9, 14]

を作る。先頭に 0 を入れるのでサイズが1つ増える。

次に (0-based インデックスで) 1 から 2 の合計は? と聞かれる。答えはひと目 3 + 4 の 7 になる。その 7 を累積和配列から取り出すには、

l = 1
r = 2
S[r + 1] - S[l]  # => 7

とする。

単純な実装 (TLE)

N, Q = gets.split.collect(&:to_i)                    # => [10, 5]
A = gets.split.collect(&:to_i)                       # => [8, 6, 9, 1, 2, 1, 10, 100, 1000, 10000]
LR = Q.times.collect { gets.split.collect(&:to_i).collect(&:pred) }
ans = LR.collect { |l, r| (l..r).sum { |i| A[i] } }  # => [15, 24, 1123, 111, 11137]
puts ans

区間の合計を1回求めるだけならこの方法でもいい。しかし何度も求める場合は効率が悪い。競プロを始めたばかりのころはここで詰んでいた。

累積和を使った実装 (AC)

187 ms
N, Q = gets.split.collect(&:to_i)                       # => [10, 5]
A = gets.split.collect(&:to_i)                          # => [8, 6, 9, 1, 2, 1, 10, 100, 1000, 10000]
LR = Q.times.collect { gets.split.collect(&:to_i).collect(&:pred) }
S = A.each_with_object([0]) { |e, m| m << m.last + e }  # => [0, 8, 14, 23, 24, 26, 27, 37, 137, 1137, 11137]
ans = LR.collect { |l, r| S[r + 1] - S[l]  }            # => [15, 24, 1123, 111, 11137]
puts ans

累積和配列の参照方法に注意

index コード
0-based S[r + 1] - S[l]
1-based S[r] - S[l - 1]

0-based に統一して毎回同じコードを書けば混乱しない。

Discussion