競プロ初中級者100問

ABC106D AtCoder Express 2

この記事で使うアルゴリズム

累積和

はじめに

カテゴリー競プロ初中級者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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import itertools
N, M, Q = map(int, input().split())
P = []
LR = [[0] * (N+1) for _ in range(N+1)] 
for m in range(M):
    l, r = map(int, input().split())
    LR[l][r] += 1

for q in range(Q):
    p = list(map(int, input().split()))
    P.append(p)

csum = []
for L in LR:
    a = list(itertools.accumulate(L))
    csum.append(a)

for p in P:
    ans = 0
    l, r = p
    for c in csum[l:r+1]:
        ans += (c[r] - c[l-1])
    print (ans)

記事情報

  • 投稿日：2020年5月4日
  • 最終更新日：2020年5月8日