AtCoder Beginner Contest 351, F : Double Sum

問題文

　https://atcoder.jp/contests/abc351/tasks/abc351_f

問題概要

　$n$ 項からなる整数列 $A = \langle A_1, A_2, \dots, A_n \rangle$ が与えられる．次式の値を求めよ．
\begin{equation*}
\sum_{ i = 1 }^{ n } \sum_{ j = i + 1 }^{ n } \max( A_j - A_i, 0 )
\end{equation*}

制約

• $2 \leq n \leq 4 \times 10^5$
• $0 \leq A_i \leq 10^8$

解法

　答えを計算するだけならば提示された数式の $\sum$ をそのまま二重 for 文に書き換えればできるわけですが，当然（？）それでは TLE するので，より高速な方法を考える必要があります．
　それにあたり，$\sum$ の値を一項ずつ足し上げて計算するのではなく，「まとめて」計算する的なことをしたい気持ちになりますが，$\max$ を取るという操作は本質的に場合分けなので，まとめて計算することとは相性が悪そうです．そこでまずは $\max$ を使わない言い換えを考えてみます．
　$\max( *, 0 )$ を噛ませなくてもよい $i, j$ というのは $A_j - A_i \geq 0$ な $i, j$ ですが，移項すれば$A_j \geq A_i$ な $i, j$ ということで，（後々の都合で）$j$ の方を基準にすれば，$A_j$ 以下の値を抜き出してできた列があると言い換えに使えそうなので，多重集合
\begin{equation*}
\{ A_i \mid i \in \{ 1, 2, \dots, j - 1 \}, A_i \leq A_j \}
\end{equation*}
を適当に順序付けて作った列を $B$ とします．$j$ を固定したときに答えに足される値は $B$ を使って
\begin{equation*}
\sum_{ i = 1 }^{ |B| }( A_j - B_i )
\end{equation*}
と書けます．このままだと使いづらいので式変形をして，
\begin{align*}
\sum_{ i = 1 }^{ |B| }( A_j - B_i )
&= \sum_{ i = 1 }^{ |B| } A_j - \sum_{ i = 1 }^{ |B| } B_i \\
&= |B| A_j - \sum_{ i = 1 }^{ |B| } B_i
\end{align*}
にします．これで，欲しい値は $B$ の要素数 $|B|$ と $B$ の総和 $\sum B$ があれば計算できることが分かりました．
　しかしながら，$B$ を毎回一から作っていては何の高速化にもならないので，もうひと工夫必要です．ところで，列上のある区間の総和を求めるのは，Segment Tree を使うことで $O( \log n )$ 時間でできます．今回の問題では区間の左端点は $1$ で固定なので，Fenwick Tree でもよい*1です．「$A_j$ 以下の値だけ抜き出した和」は難しそうですが，処理順を工夫して次のようにすれば求められます．

1. 要素数 $n$ で，$C = \{ C_1 = 0, C_2 = 0, \dots, C_n = 0 \}$ なる列 $C$ を用意する．
2. $A_j$ の昇順に $j$ に関するループを回し，各 $j$ について以下をする．
   1. $C_j \leftarrow A_j$

これで，欲しかった $\sum B$ は
\begin{equation*}
\sum_{ i = 1 }^{ j - 1 } C_i
\end{equation*}
となるので，$C$ を Segment Tree or Fenwick Tree に載せれば十分高速に求められます．$|B|$ の方も，同じく $0$ 埋めされた列を用意して，処理済みの位置に $1$ を代入していけば，同じ区間の和を取ることで求められます．ということで，それぞれの列を管理するための Segment Tree or Fenwick Tree をそれぞれ用意すれば，欲しい値をどちらも $O( \log n )$ 時間で求められます．
　全体の計算量については，まず $j$ のループ順を決定するためのソートに $O( n \log n )$ 時間がかかります．Segment Tree or Fenwick Tree については初期化に $\Theta( n )$ 時間かかり，その後，$O( \log n )$ 時間のクエリを $\Theta( n )$ 回実行するので $O( n \log n )$ 時間かかります．まとめれば全体でも $O( n \log n )$ 時間となり，TL に間に合います．

コード

// Fenwick Tree ( Binary Indexed Tree )
template < typename T >
class FenwickTree 
{
	// 中身省略
};
// FenwickTree( array size )
// int sum( 0-indexed pos )
// int sum( [ l, r ) )
// int add( pos, value )
// int change( pos, value )
// int lower_bound( value )

int main()
{
	IN( int, N );
	VI A( N );
	cin >> A;

	VI indices( N );
	iota( ALL( indices ), 0 );
	sort( ALL( indices ), [&]( const int i, const int j ){
			return A[i] == A[j] ? i < j : A[i] < A[j];
			} );

	LL res = 0;
	FenwickTree< LL > inserted( N ), values( N );
	FOR( i, indices )
	{
		res += A[i] * inserted.sum( 0, i ) - values.sum( 0, i );
		inserted.add( i, 1 );
		values.add( i, A[i] );
	}

	cout << res << endl;

	return 0;
}