Codeformula 2014 本選, D : 映画の連続視聴

問題文

　http://code-formula-2014-final.contest.atcoder.jp/tasks/code_formula_2014_final_d

問題概要

　 種類の映画がある。 番の映画の種類は で、時刻 に始まって に終わる。
　高橋くんは同じ種類の映画を連続して観ることで、より多くの幸福度を得る。同じ種類の映画を連続して 回観たとき、得られる幸福度は である。
　二つの映画 について、 であるときに限り両方の映画を観ることができる。高橋くんが得ることのできる幸福度の総和の最大値を求めよ。
　

解法

　最後に観た映画と、連続して視聴している回数をキーとする動的計画法は割とすぐに思い付きますが、これでは ある各状態から次に観る映画を 通り試す必要があり、全体では になるので TLE します。より高速な解法を考える必要があります。
　さて、上述の DP に於いて最後に観た映画を覚えていなければならないのは、次の映画を決めたときに

• 次の映画を視聴可能か
• 次の映画が直前の映画と同種か

の 2 点を判定するためです。同種映画の連続試聴回数を覚えているのは「最終的に何回連続して視聴するか」を計算するためだと考えると、ある状態からの遷移を「次に観る映画」の総当りではなく、「次に観る映画の種類と連続視聴回数」の総当りをするようにすれば、2 つ目の次元は消えます。また、この方法ならば直前の映画に関する情報の内必要なものは終了時刻だけになるので、次のような DP を考えることができます。

• dp[ 最後に観た映画の終了時刻が i ] := 得られる幸福度の和の最大値

時刻として意味のある値は高々 通りであることに注意すると、時刻を座標圧縮することで状態数は になります。なので DP のキーは「最後に観た映画の終了時刻に対応する（時刻一覧に於ける）インデックス」です。すなわち、時刻の一覧を とすると、

• dp[ i := 最後に観た映画の終了時刻が ] := 得られる幸福度の和の最大値

となります。
　状態遷移は次に視聴する映画の種類と本数を総当りしているので、ある遷移で得られる幸福度は決定しています。得られる幸福度が同じであれば（後の選択の自由度が高まるという意味で）終了時刻は早い方がよいので、「時刻 以降に種類 の映画を 回連続して観たときの終了時刻の最小値」が必要になります。映画を予め種類毎に分類しておけば、各種類に対して有効な各 に対し、時刻 以降に始まる映画を 本観たときの終了時刻の最小値を求める、と言い換えられます。これは区間スケジューリングなので、映画を終了時間の昇順にソートして、手前から順に選べるものを貪欲に採用していけばよいことになります。この計算は同じ種類の映画の本数に対し線形時間でできます。各種類の映画の本数の和は になるので、ある状態からの遷移の総当りは 時間で完了します。
　状態数が であったのでこの DP の全体では となります。DP の前処理で映画を種類ごとにソートしていますが、これは で実行できるので全体のオーダーも です。

コード

using VI = vector<int>; using VVI = vector<VI>;
using PII = pair<int,int>; using VPII = vector< pair<int,int> >;
template < typename T = int > using VT = vector<T>;
template < typename T = int > using LIM = numeric_limits<T>;

template < typename T > inline void input_n( VT< T > &out ) { copy_n( istream_iterator< T >( cin ), out.size(), out.begin() ); };

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define FOR( e, c ) for ( auto &e : c )
#define ALL( c ) (c).begin(), (c).end()

#define PB( n ) push_back( n )
#define EB emplace_back

#define fst first
#define snd second

constexpr int INF = LIM<>::max() / 2;

int main()
{
	cin.tie( 0 );
	ios::sync_with_stdio( false );

	int n;
	cin >> n;

	VI hs( n );
	input_n( hs );

	VI csum_hs;
	partial_sum( ALL( hs ), back_inserter( csum_hs ) );

	VT< VPII > movies( n );
	VI times( 1, 0 );
	REP( i, 0, n )
	{
		int m, s, e;
		cin >> m >> s >> e;

		movies[ m - 1 ].EB( s, e );

		times.PB( e );
	}
	FOR( row, movies )
	{
		sort( ALL( row ), []( const PII &a, const PII &b ){ return a.snd < b.snd; } );
	}
	
	sort( ALL( times ) );
	times.erase( unique( ALL( times ) ), times.end() );

	const int M = times.size();

	unordered_map<int,int> time2idx;
	REP( i, 0, M )
	{
		time2idx[ times[i] ] = i;
	}

	VI dp( M, -INF );
	dp[0] = 0;

	REP( i, 0, M )
	{
		REP( j, 0, n )
		{
			int t = times[i];
			VI ends;

			FOR( movie, movies[j] )
			{
				if ( t <= movie.fst )
				{
					t = movie.snd;
					ends.PB( t );
				}
			}

			REP( k, 0, ends.size() )
			{
				int &next = dp[ time2idx[ ends[k] ] ];
				next = max( next, dp[i] + csum_hs[k] );
			}
		}
	}

	cout << *max_element( ALL( dp ) ) << endl;

	return 0;
}