AtCoder Beginner Contest 186, D : Sum of difference

問題文

　https://atcoder.jp/contests/abc186/tasks/abc186_d

問題概要

　$N$ 項からなる正整数列 $\{ A_i \}_{ i \in [ 0, N ) }$ が与えられる．$$\sum_{ i = 1 }^{ N - 1 } \sum_{ j = i + 1 }^N | A_i - A_j |$$ を求めよ．

制約

• $2 \leq N \leq 2 \times 10^5$
• $| A_i | \leq 10^8$
• $A_i \in \mathbb Z$

解法

　当たり前ですが，$| A_i - A_j | = | A_j - A_i |$ です．何が言いたいかというと，差の絶対値をとるという操作は，要素の順序，即ち添字を入れ替えても結果が変わりません．よって，$A$ の並び順を適当に弄ってしまっても大丈夫ということになります．例えば扱いやすそうな形として $A$ を予めソートすることで，求めるべきものは $$\sum_{ i = 1 }^{ N - 1 } \sum_{ j = i + 1 }^N ( A_j - A_i )$$ という絶対値が外れた形になります．
　ここで，$j$ を先に固定することにして $\sum$ を書き換えると，$$\sum_{ j = 2 }^N \sum_{ i = 1 }^{ j - 1 } ( A_j - A_i )$$ となります．$j$ を固定するので外側の $\sum$ は一旦見えないことにすると $$\sum_{ i = 1 }^{ j - 1 } ( A_j - A_i )$$ となります．更に $\sum$ の中身をバラして，$$\sum_{ i = 1 }^{ j - 1 }A_j - \sum_{ i = 1 }^{ j - 1 } A_i$$ とできます．ここで一つ目の $\sum$ の中身に $i$ に関係する項は無いので $i$ の回る範囲の個数を $A_i$ にかければ値を求めることができて，$$( j - 1 )A_j - \sum_{ i = 1 }^{ j - 1 } A_i$$ となります．で，二つ目の $\sum$ は累積和そのものです．$A$ の先頭 $i$ 項の和に $S( i ) = \sum_{ k = 1 }^i A_k$ と名前をつけることにすれば，各 $j$ に対して求めたいものは $$( j - 1 )A_j - S( j - 1 )$$ です．ここで外側の $\sum$ を思い出してみれば，$$\sum_{ j = 2 }^N ( ( j - 1 ) A_j - S( j - 1 ) )$$ を得ます．ここで $\sum$ を分配すると $$\sum_{ j = 2 }^N ( j - 1 ) A_j - \sum_{ j = 2 }^N S( j - 1 )$$ となります．二項目の添字をちょっとずらすと $$\sum_{ j = 2 }^N ( j - 1 ) A_j - \sum_{ j = 1 }^{ N - 1 } S( j )$$ です．それぞれの項が何を言っているのかを読み下すと，

• 一項目 : 各 $j$ について $( j - 1 )$ と $A_j$ の積をとったものの和
• 二項目 : $A$ の累積和から末尾を取り除いたものの和

です．
　一項目は $j$ を順に動かしながら値を計算して和をとることで $\Theta( N )$ 時間で求められます．二項目も累積和を std::partial_sum するなりループで作るなりどうとでもできて，和をとることと含めて全体で $\Theta( N )$ 時間です．一番時間がかかるのは $A$ のソートで，全体では $O( N \log N )$ 時間です．

コード

main = do
	getLine
	as <- sort <$> readIntegers
	print $ sum ( zipWith (*) [ 0 .. ] as ) - sum ( init $ scanl1 (+) as )

int main()
{
	IN( int, N );
	VT< LL > A( N );
	cin >> A;

	sort( ALL( A ) );

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

	LL res = 0;
	REP( j, N )
	{
		res += j * A[j];
	}
	res -= accumulate( begin( psum ), end( psum ) - 1, 0LL );
	cout << res << endl;

	return 0;
}