Kazun の競プロ記録

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

AtCoder Beginner Contest 233 D 問題 Count Interval

AtCoder ABC ABC233 D問題 400 pts

問題

提出解答

問題の概要

次を満たすような整数の組 $(l,r)$ の個数を求めよ.

$1 \leq l \leq r \leq N$
$\displaystyle \sum_{i=l}^r A_i=K$

制約

$1 \leq N \leq 2 \times 10^5$
$|A_i| \leq 10^9$
$|K| \leq 10^{15}$

解法

$\displaystyle \sum_{i=l}^r A_i=\sum_{i=1}^{r} A_i-\sum_{i=1}^{l-1} A_i$

であるから, $k=0,1, \dots, N$ に対して,

$\displaystyle S_k=\sum_{i=1}^k A_i$

とすると, 求めるべきは $S_r-S_{l-1}=K$ となる $l \leq r$ の個数になる.

$l':=l-1$ とおくと, $S_r-S_{l'}=K, 0 \leq l' \lt r \leq N$ となる組の個数を求めることになる.

よって, 以下のようにして正解できる.

任意の整数 $x$ に対して, $D_x \gets 0$ とする.
$X \gets 0$
$D_0 \gets 1$ とする ( $S_0=0$ の分).
に対して, 以下を行う.
- $X$ に $D_{S_i-K}$ を加算する.
- $D_{S_i}$ に $1$ を加算する.
$X$ が答え.

計算量は辞書を利用して, $O(N) (O(N \log N)?)$ である.