ACL Beginner Contest, D : Flat Subsequence

問題文

　https://atcoder.jp/contests/abl/tasks/abl_d

問題概要

　数列 $\{ A_i \}$ と正整数 $K$ が与えられる．以下の条件を満たす数列 $B$ の長さとしてあり得る最大値を求めよ．

• $B$ は $A$ の部分列
• （意味のある）任意の $i$ について $| B_i - B_{ i + 1 } | \leq K$

制約

• $1 \leq N \leq 300{,}000$
• $0 \leq A_i, K \leq 300{,}000$

解法

　次のように状態をとる動的計画法を考えます

• $\mathit{ dp }[i][j] := A \text{ の } i \text{ 項目まで見たとき，末尾が } j \text{ となる } B \text{ の長さの最大値 }$

このとき，$A_i$ を使う場合の $\mathit{ dp }[i][j]$ の更新は，直前の値としてあり得るものをすべて試す形となり，$\mathit{ mi } = \max( 0, j - K ), \mathit{ ma } = \min( 300{,}000, j + K )$ として，$$\mathit{ dp }[i][j] \leftarrow \max_{ \mathit{ mi } \leq k \leq \mathit{ ma } } \mathit{ dp }[ i - 1 ][k] + 1$$ です．この DP は $O( NK )$ 時間となりますが，次のように高速化できます．
　DP 配列の $i$ の次元を消去して配列を使い回すことにします．また，更新時の更新元は範囲からの $\max$ をとっているので，Range Maximum Query を実装した Segment Tree があれば $O( \log N )$ 時間で更新時に代入するべき値を計算できます（更新も $O( \log N )$ 時間です）．よって，全体で $O( N \log N )$ 時間となって間に合います．

コード

#include <atcoder/segtree>
using namespace atcoder;

using S = int;

S op( const S a, const S b )
{
	return max( a, b );
}

S e()
{
	return 0;
}

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

	segtree< S, op, e > rmq( 300'001 );
	FOR( a, A )
	{
		const int mi = max( 0, a - K );
		const int ma = min( 300'000, a + K );
		rmq.set( a, rmq.prod( mi, ma + 1 ) + 1 );
	}

	cout << rmq.all_prod() << endl;

	return 0;
}