torus711 のアレ

主に競技プログラミングの問題について書きます.PC 以外だと数式が表示されないかもしれないです

AtCoder Beginner Contest 439, F : Beautiful Kadomatsu

問題概要

 サイズ $n$ の順列 $P = \langle P_1, P_2, \dots, P_n \rangle$ が与えられる.$P$ の部分列 $\langle A_1, A_2, \dots, A_k \rangle$ について,次の条件を満たすものの個数を(${} \bmod {998{,}244{,}353}$ で)求めよ.

  • $x = \#\{ i \in \{ 1, 2, \dots, n \} \mid 1 < i < k \land A_{ i - 1 } < A_i > A_{ i + 1 } \}$ とする
  • $y = \#\{ i \in \{ 1, 2, \dots, n \} \mid 1 < i < k \land A_{ i - 1 } > A_i < A_{ i + 1 } \}$ とする
  • $x > y$ である

なお,集合 $S$ に対し,$\#S := |S|$ である.

制約

  • $1 \leq n \leq 3 \times 10^5$

解法

 まず,$x, y$ がそれぞれどういう意味かを整理しますが,$P$ を先頭から走査したときに上昇から下降に切り替わる回数が $x$,下降から上昇に切り替わる回数が $y$ です.$P$ がサイズ $n$ の順列であることから隣接する要素が等しいことは無いので,$P$ の任意の部分列は上昇・下降を交互に繰り返すことになります.このとき,上昇から下降に転ずるには上昇から切り返すしかなく,下降から上昇に転ずるには下降から切り返すしかありません.つまり,上昇から始めたなら $x \geq y$ であり,下降から始めたなら $x \leq y$ となります.また,等号が成立するのは上昇からの切り替えしと下降からの切り返しが $1:1$ 対応するときです.従って,$x > y$ にするためには上昇から初めて下降途中で終わらせるしかありません.なので,必要条件として
\[
A_1 < A_2 \land A_{ k - 1 } > A_k
\]
があることが分かります.更に,$\langle A_1, A_2, \dots, A_{ k - 1 }, A_k \rangle$ は上昇する二項から始まって下降する二項で終わる列であるため,この部分が単峰ならば $1 = x > y = 0$ です.また,どこかで下降から上昇に転ずるならば(切り返しそれぞれに対応して)$A_{ k - 1 } > A_k$ であることから $x, y$ ともに $1$ ずつ増えることになるので.$x > y$ を維持します.従って,$A_2, A_{ 2 + 1 }, \dots, A_{ k - 1 }$ がどのような形をしていても問題文の条件を満たします.つまり,この部分の各項に対してそれを選ぶ・選ばないの二択を自由に選べるので $2$ 冪の選択肢があります.
 ここまでをまとめると,各 $i \in \{ 1, 2, \dots, n \}$ に対して
\begin{align*}
L_i &= \#\{ j \in \{ 1, 2, \dots, i - 1 \} \mid A_j < A_i \} \\
R_i &= \#\{ j \in \{ i + 1, i + 2, \dots, n \} \mid A_i < A_j \}
\end{align*}
とすれば,$P_i$ を部分列の先頭の次の要素とする方法が $L_i$ 通り,『部分列の末尾の手前の要素とする方法が $R_i$ 通りとなります.また,$P$ の $l$ 項目を部分列の先頭の次,$r$ 項目を末尾の手前としたとき,$l + 1$ 項目から $r - 1$ 項目の選び方は $2^{ r - l - 1 }$ 通りとなります.
 ということで $L, R$ を計算したい気持ちになりますが,これは Segment TreeFenwick Tree を用いて $O( n \log n )$ 時間で構築できます.具体的には,区間和を求められるようにして各 $i$ に対して

  1. 半開区間 $[ 0, i )$ の和を求める
  2. 位置 $P_i$ に $1$ を加算する

ということをすると.$R$ は先頭からの走査,$L$ は逆向きの走査で計算できます.
 答えの計算は $k = 3$ か否かで分けて考えます.$k = 3$ のとき,答えに対する寄与は各 $i \in \{ 1, 2, \dots, n \}$ で
\[
L_i \times R_i
\]
です.そうでない場合,$i < j$ なる $i, j$ をとると
\[
L_i \times R_j \times 2^{ j - i - 1 }
\]
が答えに寄与します.$j$ を右に動かして走査する気持ちになると,$i < j$ なる $i$ を末尾位置とする部分列を考えたとき,

  • 空でない部分列に対して $P_j$ をくっつけるか否かを選べるので個数が $2$ 倍になる
  • $P_j$ のみからなる部分列を新たに作れるので $L_j$ 個増える

という選択肢があって,その時点での個数を $t$ とすれば*1
\[
2t + L_i
\]
となります.
 従って,$L, R$ を Segment Tree 等を用いて計算してから答えを計算することで,$O( n \log n )$ 時間で問題を解くことができます.

コード

#include <atcoder/modint>
using namespace atcoder;
using MINT = modint998244353;
inline ostream& operator<<( ostream &s, const MINT &t ){ s << t.val(); return s; }

// 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 P( N );
	cin >> P;
	MAP_PRED( P );

	VT< MINT > left( N ), right( N );
	FenwickTree< MINT > fenwick1( N ), fenwick2( N );
	REP( i, N )
	{
		fenwick1.add( P[i], 1 );
		left[i] = fenwick1.sum( 0, P[i] );
	}
	for ( int i = N - 1; 0 <= i; --i )
	{
		fenwick2.add( P[i], 1 );
		right[i] = fenwick2.sum( 0, P[i] );
	}

	MINT res = 0;
	REP( i, N )
	{
		res += left[i] * right[i];
	}
	MINT tails = 0;
	REP( i, N )
	{
		res += tails * right[i];
		tails = tails * 2 + left[i];
	}

	cout << res << endl;

	return 0;
}

*1:"tails" から