AtCoder Beginner Contest 346, C : Σ

問題文

　https://atcoder.jp/contests/abc346/tasks/abc346_c

問題概要

　$n$ 項からなる正整数列 $A = \langle A_1, A_2, \dots, A_n \rangle$ と正整数 $k$ が与えられる．
　$1$ 以上 $k$ 以下の整数の内，$A$ に含まれないものの和はいくらか？

制約

• $1 \leq n \leq 2 \times 10^5$
• $1 \leq k \leq 2 \times 10^9$
• $1 \leq A_i \leq 2 \times 10^9$

解法

　まず思いつくのは $1$ 以上 $k$ 以下の整数それぞれについて $A$ に含まれるかどうかを判定して条件を満たすものについて和を取ることですが，$k$ の制約から，$\Omega( k )$ 時間かけると TLE するので，これはうまくいきません．別の方法を考える必要があります．
　答えとなる値についてちょっとよく考えてみます．「$1$ 以上 $k$ 以下の整数の内 $A$ に含まれるものの和」を $x$ とすれば，「$1$ 以上 $k$ 以下の整数の和」ヵら $x$ を引いた値が欲しい値です．「$1$ から $k$ の和」を求めるには公式があって，
\begin{equation*}
\frac{ k ( k + 1 ) }{ 2 }
\end{equation*}
で求まります．$x$ については，$A$ から

• $k$ より大きい値を取り除く
• 重複を取り除く

の 2 つのことをして得られる集合を
\begin{equation*}
\hat A = \{ a \in A \mid a \leq k \}
\end{equation*}
とすれば，
\begin{equation*}
\sum \hat A
\end{equation*}
となります．このそれぞれを実現する方法は言語によるかと思いますが，C++ ではそれぞれ，

• std::remove_if
• std::sort + std::unique

で実現できます．
　$x$ の計算は $O( 1 )$ 回の正数演算で計算できるので，$O( 1 )$ 時間で求まります．$\hat A$ の計算はソートする部分が一番重く，$O( n \log n )$ 時間かかりますが，十分に高速です．
　必要な値を求める方法が分かったので，
\begin{equation*}
\frac{ k ( k + 1 ) }{ 2 } - \sum \hat A
\end{equation*}
を出力することで問題に正答できます．

コード

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

	sort( ALL( A ) );
	A.erase( remove_if( ALL( A ), bind( greater< int >(), _1, K ) ), end( A ) );
	A.erase( unique( ALL( A ) ), end( A ) );

	cout << K * ( K + 1 ) / 2 - accumulate( ALL( A ), 0LL ) << endl;

	return 0;
}