Kazun の競プロ記録

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

AtCoder Beginner Contest 281 E問題 Least Elements

AtCoder ABC281 ABC E問題 500 pts

問題

提出解答

問題の概要

$N$ 個の整数 $A_1, \dots, A_N$ が与えられる. $q=1,2, \dots, N-M+1$ について以下の問題を解け.

$M$ 個の整数 $A_q, A_{q+1}, \dots, A_{q+M-1}$ を昇順に並べたとき, 先頭 $K$ 個の総和を求めよ.

制約

$1 \leq K \leq M \leq N \leq 2 \times 10^5$
$1 \leq A_i \leq 10^9$

解法

$(A_q, \dots, A_{q+M-1})$ の昇順 $K$ 番以内と $(A_{q+1}, \dots, A_{q+M})$ の昇順 $K$ 番以内に選ばれる整数の違いは高々2個である. このことを利用し, $q$ の答えから $(q+1)$ の答えを高速に求めることにする.

次のようにして求める.

$q$ の時の解答を $X_q$ とする.
$X_1$ をソートを利用して求める.
多重集合 $E,F$ を用意する. この $E$ は昇順 $K$ 番以内, $F$ はそれ以外の要素を管理する要素である.
の順に以下を行う.
- $A_q \leq \max E$ ならば, $E$ から $A_q$ を削除し, そうでないならば $F$ から $A_q$ を削除する.
- $A_{q+M} \leq \max E$ ならば, $E$ に $A_{q+M}$ を挿入し, そうでないならば $F$ から $A_{q+M}$ を追加する.
- $|E| \neq K$ ならば, 「 $E$ の最大値を $F$ に移動させる」または「 $F$ の最小値を $E$ に移動させる」を行い, $|E|=K$ になるようにする (この操作はどちらか1回のみで済む).
- $X_{q+1}$ は $X_q$ に $E$ の差分を加算した値である.

順序つき集合を利用することによって, $X_1, \dots, X_{N-M+1}$ を $O(N \log N)$ 時間で求めることができる.