Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 256 F問題 Cumulative Cumulative Cumulative Sum

AtCoder ABC ABC256 F問題 500 pts

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

長さ N の整数列 A=(A_1, \dots, A_N) が与えられる. 次の Q 個のクエリを順に処理せよ.

  • Type 1: A_x \gets v
  • Type 2: B を A の累積和, C を B の累積和, D を C の累積和としたとき, D の第 x 項目を求めよ.

制約

  • 1 \leq N \leq 2 \times 10^5
  • 1 \leq Q \leq 2 \times 10^5
  • 0 \leq A_i \leq 10^9
  • 1 \leq x \leq N
  • 0 \leq v \leq 10^9

解法

A の第 i 項目, B の第 j 項目, C の第 k 項目, D の第 l 項目をそれぞれ A_i, B_j, C_k, D_l と書くことにする. このとき,

  • \displaystyle B_j=\sum_{i=1}^j A_i
  • \displaystyle C_k=\sum_{j=1}^k B_j=\sum_{j=1}^k \sum_{i=1}^j A_i=\sum_{i=1}^k \left(\sum_{j=i}^k 1 \right) A_i=\sum_{i=1}^k (k+1-i) A_i
  • \displaystyle D_l=\sum_{k=1}^l C_k=\sum_{k=1}^l \sum_{i=1}^k (k+1-i) A_i=\sum_{i=1}^l \left(\sum_{k=i}^l (k+1-i) \right) A_i
    \displaystyle =\sum_{i=1}^l \dfrac{1}{2} (l+2-i)(l+1-i) A_i

である. このとき, \dfrac{1}{2} (l+2-i)(l+1-i) について, l の式と i の式の積和に変形できないかどうかを考える. すると,

\begin{align} \dfrac{1}{2} (l+2-i)(l+1-i) &=\dfrac{1}{2} i(i+1)+\dfrac{1}{2} (l+2)(l+1-2i) \\ &=\dfrac{1}{2} i(i+1)-i(l+2)+\dfrac{1}{2} (l+1)(l+2) \end{align}

であるから, 列 E_0, E_1, E_2 をそれぞれ

\displaystyle (E_0)_p:=\sum_{i=1}^p A_i, \quad (E_1)_q:=\sum_{i=1}^q i A_i,\quad (E_2)_r:=\sum_{i=1}^r \dfrac{1}{2} i(i+1) A_i

とすると,

\begin{align} D_l &=\sum_{i=1}^l \dfrac{1}{2} i(i+1) A_i-(l+2) \sum_{i=1}^l i A_i+\dfrac{1}{2} (l+1)(l+2) \sum_{i=1}^l A_i \\ &=(E_2)_l-(l+2) (E_1)_l+\dfrac{1}{2} (l+1)(l+2) (E_0)_l \end{align}

となる.

ここで, E_0, E_1, E_2 の管理と計算方法について, 各点更新と区間和が必要であり, 区間和は整数の和であることから, Binary Indexed Tree というデータ構造を利用することによって, 各クエリを Type 1, Type 2 共に O(\log N) で処理, 計算できるので, 全体で O(Q \log N) 時間で Q クエリまとめて処理できた.