AtCoder, エイシングプログラミングコンテスト 2022 (AtCoder Beginner Contest 255), D : ±1 Operation 2

問題文

　https://atcoder.jp/contests/abc255/tasks/abc255_d

問題概要

　項数 $N$ の数列 $\{ A_i \}_{ i \in [ 0, N ) }$ が与えられる．この数列を以下のように変更することを「操作」と呼ぶ．

1. $0 \leq i < N$ なる整数 $i$ を選ぶ（プログラム的にアレなので 0-based に変えています）
2. $A_i$ を $A_i + 1$ または $A_i - 1$ のいずれかで置き換える

　以下で説明されるクエリが $Q$ 個与えられるので，それぞれ答えよ．

1. 整数 $X$ が（クエリ毎に独立に）与えられる
2. 数列 $A$ の要素が全て $X$ と等しくなるようにするために必要な操作回数の最小値を出力せよ

制約

• $1 \leq N, Q \leq 2 \times 10^5$
• $0 \leq A_i, X \leq 10^9$

解法

　まず分かることとして，$X$ が与えられたとき，各 $A_i$ を $X$ から遠ざかる方向に変化させることには損しかないので，必ず $X$ との差（の絶対値）が小さくなるように操作することになります．そうしたとき，操作回数の総和は $X$ との差（の絶対値）の和，$$\sum_{ i = 0 }^{ N - 1 }| X - A_i |$$ となります．絶対値が付いていると扱いづらいので，これを外します．「$X$ に近づけるように操作する」を言い換えて，

1. $A_i \leq X$ なら（必要に応じて） $A_i + 1$ で置き換える
2. $A_i > X$ なら $A_i - 1$ で置き換える

と分解します．そのときの操作回数は，

1. $A_i \leq X$ な $i$ についての $X - A_i$ の総和
2. $A_i > X$ な $i$ についての $A_i - X$ の総和

の和となります．式で書くなら，ええっと…… $$\sum\{ X - A_i \mid i \in [ 0, N ), A_i \leq X \} + \sum\{ A_i - X \mid i \in [ 0, N ), X < A_i \}$$ みたいな……？
　これで求めるべき値が定まりはしましたが，クエリの度に $A$ の全要素を舐めるみたいなことをすると TLE してしまうので，計算方法を工夫する必要があります．なんとなくやりたいこととしては，毎回全部を舐めてはいられないので，なんとかして「まとめて」処理したい気がします．いきなり言い出すとやや天啓かもしれませんが，よく言われる ^{[誰によって？]}考え方として，「順序が関係ないものはソートすると嬉しくなるか考えよう」というのがあります．$A$ については各要素の $X$ との差だけに興味があって，中身の順序には興味がないのでソートしてしまえます．そうすると，加算の操作は「ある位置」より手前に，減算の操作は「ある位置」以降にある要素を対象とすることが分かります．また，「ある位置」は二分探索（C++ では std::lower_bound とか）によって $O( \log N )$ 時間で求めることができます．なお，「ある位置」をちゃんと言うと「$X \leq A_i$ となる最初の要素」です．
　$A$ が昇順ソート済み，「ある位置」を指す添字を $p$ としたとき，求めたい値は式で書くと $$\sum_{ i = 0 }^{ p - 1 }( X - A_i ) + \sum_{ i = p }^{ N - 1 }( A_i - X )$$ と，一個前のものよりすっきりします．どういう表記を認めるのかとか，その辺はちょっとあるのですが……．とはいえ，$A_i$ と $X$ との大小関係にる場合分けが消えるので，扱いやすくはなっています．更に，各項の $X$ というのは $\sum$ の添字が何回回るのかを考えると外に出すことができて，$$pX - \sum_{ i = 0 }^{ p - 1 } A_i + \sum_{ i = p }^{ N - 1 }A_i - ( N - p ) X$$ と変形できます．
　こうなると残っている難しい部分は $\sum_{ i = 0 }^{ p - 1 }A_i$ と $\sum_{ i = p }^{ N - 1 }A_i$ で，これは日本語で言い直すと「配列上の区間の和」です．これを求める方法として，「累積和」と呼ばれるものがあります．この記事が既に長いので検索ワードとして名前だけ紹介して詳細は割愛しますが，このテクニックを使うと，長さ $N$ の配列に対して

• $\Theta( N )$ 時間の前計算
• $O( 1 )$ 時間の区間和計算

を実現できます．
　従いまして，まとめるとこの問題は

1. 与えられた $A$ をソートして新たな $A$ とする
2. $A$ の累積和を計算しておく
3. クエリ毎に以下をする
   1. 加算操作の対象となる要素と減算操作の対象となる要素の境界を二分探索する
   2. 各範囲の総和を，累積和の情報を使って求める
   3. 得られた値たちでかけたり足したりして答えを計算する

といった手順で解くことができます．

計算量について

　上記の解法の各手順での計算量をまとめると，

• $A$ のソートに $O( N \log N )$ 時間
• 累積和の計算に $\Theta( N )$ 時間
• $Q$ 回のクエリについて，
  • 二分探索に $O( \log N )$ 時間
  • 答えの計算に $O( 1 )$ 時間

がかかります．よって全体では $O( N \log N + Q \log N ) = O( ( N + Q ) \log N )$ 時間となります．

コード

int main()
{
	cin.tie( nullptr );
	ios::sync_with_stdio( false );
	cout << setprecision( 12 ) << fixed;

	IN( int, N, Q );
	VT< LL > A( N );
	cin >> A;
	sort( ALL( A ) );

	VT< LL > psum( 1 );
	partial_sum( ALL( A ), BI( psum ) );

	REP( Q )
	{
		IN( LL, X );

		const int pos = lower_bound( ALL( A ), X ) - begin( A );
		const LL L = pos * X - psum[ pos ];
		const LL R = psum[N] - psum[ pos ] -  ( N - pos ) * X;

		cout << L + R << '\n';
	}

	cout << flush;

	return 0;
}