ABC480Dを解いた【累積和】

筆者はレート800前後の茶～緑コーダ

ABC480のD問題を解いていく

実装コード

二段階の累積和と三角数を使い、A の先頭から長さ 1, 2, …, N の列を順につなげた数列の区間和を求める。

• N、Q と数列 A を受け取る
• S[i] を A の先頭 i 個の和として、累積和 S を作る
• T[i] を S[1] から S[i] までの和として、先頭 i ブロック分の合計を前計算する
• f(M) では、先頭 M 個に完全に含まれるブロック数 i を (isqrt(8 * M + 1) - 1) // 2 で求める。isqrt を使うことで浮動小数点の誤差を避ける
• j = M - i * (i + 1) // 2 で次のブロックに含まれる要素数を求め、T[i] + S[j] を先頭 M 個の和として返す
• 各クエリの L, R に対し、f(R) - f(L - 1) で区間和を求めて出力する

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
from math import isqrt
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

def main():
    N, Q = rLI()
    A = rLI()

    S = [0] * (N + 1)
    T = [0] * (N + 1)
    for i in range(N):
        S[i + 1] = S[i] + A[i]
        T[i + 1] = T[i] + S[i + 1]

    def f(M):
        i = (isqrt(8 * M + 1) - 1) // 2
        j = M - i * (i + 1) // 2
        return T[i] + S[j]

    for _ in range(Q):
        L, R = rLI()
        ans = f(R) - f(L - 1)
        print(ans)

if __name__ == '__main__':
    main()

感想

累積和をもう一段取ると、ブロック単位の和をまとめて求められるのがポイントだった。
三角数から完全に含まれるブロック数と残りの要素数に分けると、数列を実際に作らずに区間和を計算できる。