Pythonで累積和を実装してみる-ABC122

累積和

最近Atcoderで頻出な感じがする累積和をPythonで実装していきます。

累積和とはその名の通り、累積の和をとる前処理をすることによってクエリの速度を上げることができるアルゴリズムになります。

実際に問題を解きながら解説していきます。

問題

今回は以下の問題を扱います。
atcoder.jp

問題文

A, C, G, T からなる長さ N の文字列 S が与えられます。以下の Q個の問いに答えてください。
問 i(1≤i≤Q): 整数 li,ri (1≤li<ri≤N)が与えられる。S の先頭から li 文字目から ri文字目までの (両端含む) 部分文字列を考える。この文字列に AC は部分文字列として何回現れるか。
文字列 Tの部分文字列とは、T の先頭と末尾から 0文字以上を取り去って得られる文字列です。
例えば、ATCODER の部分文字列には TCO, AT, CODER, ATCODER, (空文字列) が含まれ、AC は含まれません。

制約

2≤N≤10^5
1≤Q≤10^5
Sは長さ Nの文字列
Sの各文字は A, C, G, T のいずれか
1≤li<ri≤N

実装

この問題は指定された範囲内にACが何回現れるかを答える問題になります。

まずは累積和を使わずに単純にそのまま指定された範囲内にACが何回現れるのかを数えていくプログラムを書いてみます。

n,q = map(int,input().split())
s = input()
for i in range(q):
    l,r = map(int,input().split())
    ans = 0
    for j in range(l-1,r-1):
        if s[j]=="A" and s[j+1]=="C":
            ans += 1
    print(ans)

入力例は通りますが、いざ提出してみると……

f:id:nashidos:20191219132503p:plain

圧倒的TLEですね。制約からもわかる通り、NとQがかなり大きいので処理を高速化する必要があります。

先ほどのコードでは毎回毎回ACの部分集合を数えていました。しかし、それでは間に合いません。

そこで累積和です。

あらかじめACの部分集合がいくつあるのかを前処理で数えておけば毎回数える必要がないので高速になります。

以下が累積和を利用したコードです。

n,q = map(int,input().split())
s = input()
ls = [0]
for i in range(n-1):
    if s[i] == "A" and s[i+1] == "C":
        ls.append(ls[i]+1)
    else:
        ls.append(ls[i])

for i in range(q):
    l,r = map(int,input().split())
    print(ls[r-1]-ls[l-1])