この記事で使うアルゴリズム
累積和
はじめに
カテゴリー競プロ初中級者100問では、Qiita
にて@e869120
さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編:目指せ水色コーダー!】としてまとめられている100問をPython
で解いています。
全問題の一覧はこちらです
問題
https://atcoder.jp/contests/abc106/tasks/abc106_d
この問題では累積和の計算が求められます。
方針
愚直にクエリ毎に全ての列車を調べると、計算量はO(MQ)
となり間に合いません。
そこでに二次元空間に展開することで計算量を削減します。具体的には縦軸を始点L
、横軸を終点R
として、 区間Li
からRi
まで走る列車i
を点として表します。全く同一の区間を走る列車もあるので、二次元配列上にカウントするのが良いでしょう。
このようにデータを展開しておくことで、クエリに対して二次元配列上の対応する長方形の中の数の和を求めれば良いことになります。この方法だと、計算量はO(N^2 Q)
となり、今回の制約では先ほどとあまり変わりません。高速化が必要です。
区間の和を複数回求める場合、累積和を使うことで高速化できる場合があります。(O(N)
がO(1)
になります。)
累積和
とはA = [a_1, a_2,..., a_n]
に対してs_i = a_1 + a_2 + ... + a_i
のように定義される値です。いかに例を示します。
配列の値 | 1 | 4 | 0 | 2 | 3 |
---|---|---|---|---|---|
累積和 | 1 | 5 | 5 | 7 | 10 |
Python
で累積和を実装する場合、itertools.accumulate
を用いると良いでしょう。
今回の場合、二次元配列ですので二次元累積和を考えることもできますが、横軸方向を累積和で、縦軸方向をループで処理すればPyPy
で十分間に合います。計算量はO(NQ)
となります。
コード
1 | import itertools |
記事情報
- 投稿日:2020年5月4日
- 最終更新日:2020年5月8日