Kazun の競プロ記録

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

AtCoder Beginner Contest 265 D問題 Iroha and Haiku (New ABC Edition)

AtCoder ABC ABC265 D問題 400 pts

問題

提出解答

問題の概要

長さ $N$ の非負整数列 $A=\left(A_i \right)_{i=0}^{N-1}$ に対して, 以下を満たす整数 $x,y,z,w$ は存在するか?

$0 \leq x \lt y \lt z \lt w \leq N$
$A_x+A_{x+1}+\dots+A_{y-1}=P$
$A_y+A_{y+1}+\dots+A_{z-1}=Q$
$A_z+A_{z+1}+\dots+A_{w-1}=R$

制約

$3 \leq N \leq 2 \times 10^5$
$1 \leq A_i \leq 10^9$
$1 \leq P,Q,R \leq 10^{15}$

解法

$S=\left(S_i \right)_{i=0}^N$ を $A$ の累積和とする. つまり,

$\displaystyle S_0:=0, \quad S_i:=\sum_{j=1}^{i-1} S_j=S_{i-1}+A_{i-1} \quad (1 \leq i)$

と定義する.

このとき,

$A_x+A_{x+1}+\dots+A_{y-1}=P \iff S_y-S_x=P$

同様にすると, 求めるべき問題を

$0 \leq x \lt y \lt z \lt w \leq N$
$S_y-S_x=P$
$S_z-S_y=Q$
$S_w-S_z=R$

と簡略化することができた.

$0 \leq x \lt N$ を1つ固定する. するとこの条件下では, 条件は

$0 \leq x \lt y \lt z \lt w \leq N$
$S_{x+1}, S_{x+2}, \dots, S_N$ の中に $S_x+P$ が存在する.
$S_{y+1}, S_{y+2}, \dots, S_N$ の中に $S_x+P+Q$ が存在する.
$S_{z+1}, S_{z+2}, \dots, S_N$ の中に $S_x+P+Q+R$ が存在する.

そして, $A$ の各項が正であり, $x \lt y \lt z \lt w$ であるから, 条件は

$S_0, S_1, \dots, S_N$ の中に $S_x+P, S_x+P+Q, S_x+P+Q+R$ が全て存在する

になる. $S_0, S_1, \dots, S_N$ は $A$ が与えられたらその瞬間固定されるので, 各項を要素に持つ集合で記録することで存在判定を高速に処理することができる.

この $x$ を全て動かすことにより, 計算量 $O(N)$ で求めることができた.