AtCoderの累積和入門：Typical90 010をPythonで解く区間和の考え方

累積和は、値が変わらない配列に対して、同じ先頭部分を何度も足す代わりに先頭から何個分までの合計を一度だけ保存する前計算です。Typical90 010では、1組用と2組用の累積配列を作れば、1始まりの閉区間[L, R]の合計を、それぞれprefix[R] - prefix[L - 1]で求められます。

この方法がそのまま使えるのは、前計算後に点数が変わらず、同じ配列へ複数の区間問い合わせを行う場合です。Typical90 010には点数の更新がないため適合します。Pythonで標準入力を読み、listとfor文を使えることを前提に、添字の意味から確認します。

同じ先頭部分を前計算して再利用する

区間ごとにL番からR番までを足し直す方法では、1件の問い合わせに最大O(N)かかります。Typical90 010ではNとQがともに最大100000なので、全問い合わせではO(NQ)になります。そこで、複数の問い合わせに重なる「先頭からの合計」を保存し、足し直しをなくします。

prefix1[i]を先頭i人のうち1組の点数合計、prefix2[i]を先頭i人のうち2組の点数合計と定義します。どちらもprefix[0] = 0から始めます。この0番目を置くことで、L = 1でも別の分岐を作らず、同じ差分式を使えます。

5人の例で2本の累積配列を追う

次のデータは、仕組みを追うために作成した例であり、公式サンプルではありません。各組の点数を、学籍番号の左から順に累積します。

学籍番号i	(C_i, P_i)	prefix1[i]	prefix2[i]
0	開始前	0	0
1	(1, 4)	4	0
2	(2, 7)	4	7
3	(1, 3)	7	7
4	(2, 5)	7	12
5	(1, 6)	13	12

入力は順に[(1, 4), (2, 7), (1, 3), (2, 5), (1, 6)]です。したがって、完成した配列はprefix1 = [0, 4, 4, 7, 7, 13]、prefix2 = [0, 0, 7, 7, 12, 12]になります。1人読むたびに直前の2つの合計を引き継ぎ、その人が属する組の配列だけへ点数を加えています。

閉区間[2, 4]を差で取り出す

prefix1[4]には学籍番号1〜4の1組合計が入り、prefix1[1]には区間の手前、学籍番号1までの1組合計が入っています。前者から後者を引けば、学籍番号2〜4だけが残ります。2組も同じです。

組	先頭からR = 4まで	先頭からL - 1 = 1まで	区間[2, 4]
1組	prefix1[4] = 7	prefix1[1] = 4	7 - 4 = 3
2組	prefix2[4] = 12	prefix2[1] = 0	12 - 0 = 12

直接見ても、学籍番号2〜4の1組は3番の3点、2組は2番の7点と4番の5点なので、結果は3 12です。差分で求めた値と一致します。R + 1へ変換していないのは、prefix[i]を「先頭i人」と定義し、問題文の1始まりの端点をそのまま使っているからです。

Typical90 010で求められる入出力

Typical90 010 – Score Sum Queriesの公式問題では、N人それぞれのクラスC_iと点数P_iが与えられます。その後のQ個の問い合わせごとに、両端を含む学籍番号L_jからR_jまでの1組と2組の合計を、この順に1行で出力します。

• 1 ≤ N ≤ 100000
• 1 ≤ C_i ≤ 2
• 0 ≤ P_i ≤ 100
• 1 ≤ Q ≤ 100000
• 1 ≤ L_j ≤ R_j ≤ N

入力値はすべて整数です。入力はN、続くN行のC_i P_i、Q、続くQ行のL_j R_jです。公式サンプル1は次の1件の問い合わせを含みます。

7
1 72
2 78
2 94
1 23
2 89
1 40
1 75
1
2 6

期待される出力は次のとおりです。

63 261

公式サンプル1を掲載コードの配列で追う

公式サンプル1の入力を掲載コードへそのまま渡すと、前計算の各時点で配列は次のようになります。表のiはコードのfor i in range(1, n + 1)の添字で、prefix1[i]とprefix2[i]はその時点までのクラス別合計です。

i	入力の(C_i, P_i)	prefix1[i]	prefix2[i]
0	開始前	0	0
1	(1, 72)	72	0
2	(2, 78)	72	78
3	(2, 94)	72	172
4	(1, 23)	95	172
5	(2, 89)	95	261
6	(1, 40)	135	261
7	(1, 75)	210	261

問い合わせの[2,6]では、掲載コードがleft = 2、right = 6を読み、prefix1[right] - prefix1[left - 1]とprefix2[right] - prefix2[left - 1]を計算します。したがって、prefix1[6]-prefix1[1]=135-72=63、prefix2[6]-prefix2[1]=261-0=261となり、出力の63 261へ対応します。これは公式サンプル1の入力から作った追跡であり、前節の5人の構成例（公式サンプルではない）とは別です。

AtCoderの公式overall editorialには問題010の公式解説への外部リンクがありますが、AtCoder内のこのページ自体には解法本文が掲載されていません。以下の式とPythonコードは、公式問題の入出力仕様と制約から組み立てた参考実装です。公式実装の転載でも、提出してACを確認したコードでもありません。

入力から出力までを含むPython参考コード

import sys


def solve() -> None:
    data = iter(map(int, sys.stdin.buffer.read().split()))
    n = next(data)

    prefix1 = [0] * (n + 1)
    prefix2 = [0] * (n + 1)

    for i in range(1, n + 1):
        c = next(data)
        p = next(data)
        prefix1[i] = prefix1[i - 1]
        prefix2[i] = prefix2[i - 1]
        if c == 1:
            prefix1[i] += p
        else:
            prefix2[i] += p

    q = next(data)
    answers = []
    for _ in range(q):
        left = next(data)
        right = next(data)
        sum1 = prefix1[right] - prefix1[left - 1]
        sum2 = prefix2[right] - prefix2[left - 1]
        answers.append(f"{sum1} {sum2}")

    print("\n".join(answers))


if __name__ == "__main__":
    solve()

prefix1とprefix2は長さN + 1です。各人について両方の直前値をコピーし、所属する組だけに点数を加えます。問い合わせでは、入力のleftとrightを1始まりのまま使い、区間の手前をleft - 1で参照します。

参考コードが正しい理由

前計算中に保つ条件は、「i人目まで読み終えた時点で、prefix1[i]は学籍番号1〜iの1組合計、prefix2[i]は同じ範囲の2組合計である」です。開始時のi = 0では、まだ誰も含まないため両方とも0で、この条件を満たします。

i人目を読むと、まずi - 1人目までの合計を両方コピーします。そのうえで、C_i = 1ならprefix1[i]だけ、C_i = 2ならprefix2[i]だけへP_iを加えます。したがって、どちらの組についても先頭i人分の合計という条件が保たれます。

問い合わせでは、prefix_c[R]に学籍番号1〜Rのクラスcの合計、prefix_c[L - 1]に学籍番号1〜L - 1の合計が入っています。共通する先頭L - 1人分を引くと学籍番号L〜Rだけが残るため、出力は求める閉区間のクラス別合計になります。

計算量

前計算ではN人を1回ずつ処理するのでO(N)です。各問い合わせは4回の配列参照と2回の減算で答えられるため1件あたりO(1)、全体ではO(N + Q)です。掲載コードはdataとして入力全体のトークンを保持し、問い合わせごとの文字列をanswersへ蓄積するため、コード全体の追加メモリはO(N + Q)です。prefix1とprefix2の2本の累積配列だけの作業領域は、どちらも長さN + 1なのでO(N)です。

公式サンプルと境界条件の確認

掲載した参考コードへ、公式問題にある4組のサンプル入力をそれぞれ無加工で与え、期待される出力全文と照合しました。独自に作った5人の例や境界入力とは分けて確認しています。

公式サンプル	照合した内容	結果
1	1件の出力63 261	期待出力と全文一致
2	10件の問い合わせ、10行の出力	期待出力と全文一致
3	N = 1、同じ問い合わせ3回、2組の合計0	3行とも100 0で全文一致
4	20件の問い合わせ、20行の出力	期待出力と全文一致

式とコードの境界は次のように対応します。

• L = 1：L - 1 = 0なので、初期値prefix_c[0] = 0を引けばよく、特別な分岐は不要です。
• L = R：prefix_c[R] - prefix_c[R - 1]となり、その1人がクラスcなら点数、別のクラスなら0が残ります。
• N = 1：公式サンプル3で、唯一の区間[1, 1]を同じ式で処理できます。
• 片方の組が区間内にいない：その組の累積値は区間の前後で増えないため、差は0です。
• P_i = 0：対応する累積値に0を加えるだけなので、定義と差分式は変わりません。
• 同じ問い合わせを繰り返す：前計算後の配列を変更しないため、毎回同じ2つの差を計算できます。

反対に、途中で点数を書き換える問い合わせが加わると、変更位置以降の累積値も変わります。Typical90 010には値の更新がないためこの実装でよい、という適用条件を区別しておきます。

次の練習はABC129 Bで累積状態を1変数にする

ABC129 B – Balanceの公式問題では、2 ≤ N ≤ 100個の重りを、1 ≤ T < Nとなる位置で左右に分け、2つの合計の差の絶対値を最小にします。Typical90 010のように任意区間へ何度も答える問題ではありません。

分割位置を左から1回だけ調べるので、過去の全prefixを配列へ保存する必要はありません。最初にtotalへ全体の合計を入れ、分割を1つ進めるたびにleft_sumへ左側へ移った重りを加えます。その時点の右側はtotal - left_sumです。最後の要素まで左側へ入れると右側が空になるため、調べるのはT < Nの範囲です。

公式サンプル1のN = 3、重さ[1, 2, 3]では、T = 1の差は|1 - 5| = 4、T = 2の差は|3 - 3| = 0で、最小値は0です。left_sumを更新してから差を評価する順序も確認できます。ABC129の公式解説PDFも、先頭側の和を保持し、全体和との差を調べるO(N)の考え方を説明しています。

この練習で確かめたい違いは、累積和を使うときに常に配列が必要なわけではないことです。後から任意の区間を参照するTypical90 010では全prefixを保存し、左から一度だけ進むABC129 Bではleft_sumとtotalだけを持ちます。後者はO(N)時間、O(1)追加メモリで進められます。