Official

E - Colorful Subsequence Editorial by mechanicalpenciI


便宜上、元の列のさらに左に色 \(0\) で 価値が \(0\) のボールが \(1\) つ置かれていて、 そのボールは取り除かれないものとします。
(このボールは左から \(0\) 番目のボールとして扱います。)

まず、単純な動的計画法による解法として次のようなものが考えられます。
\(0\leq i\leq N\), \(0\leq j\leq K\), \(0\leq k\leq N\) について、\(DP[i][j][k]\) を、「左から \(i\) 番目までのボールまでが存在するとし、そのうち \(j\) 個を取り除く方法であって、残ったボールを並べた時に同じ色のボールが連続せず、かつ一番右のボールが色 \(k\) であるようなものについての、残ったボールの価値の総和の最大値」として定義します。ただしそのような取り除き方が存在しない場合( \(i<j\) の場合を含む)は \(-\infty\) とします。
最終的な答えは \(\displaystyle \max_{0\leq k\leq N} dp[N][K][k]\) です。

まず、\(DP[0][0][0]=0\) であり、\((j,k)\neq (0,0)\) について\(DP[0][j][k]=-\infty\) です。次に、ある \(i=i_0\) についての \(DP[i][j][k]\) の値について、\(k\neq C_{i_0}\)のとき、\(i_0\) 番目のボールは取り除くほかなく、\(DP[i_0][j][k]=DP[i_0-1][j-1][k]\)\(j=0\) のときは \(-\infty\))となります。\(k=C_{i_0}\) のとき、\(i_0\) 番目のボールは取り除く場合の最大値は先と同じく \(DP[i_0-1][j-1][k]\)\(j=0\) のときは \(-\infty\) ), 取り除かない場合は \(\displaystyle \max_{k’\neq k} DP [i_0-1][j][k’]+V_{i_0}\) であり、この \(2\) つのうち大きい方の値となります。
しかし、これを愚直に実装すると \(O(N^3K)\) または \(O(N^2K)\)の計算量となり間に合いません。

ここで、遷移式を観察すると、各 \((i,j)\) について、\(k\)\(DP[i][j][k]\) のペア \((k, [i][j][k])\) を、\(DP[i][j][k]\) の値が大きい順に \(2\) つ記録しておけば良い事が分かります。
ここで、\(DP[i][j][k]\) の値が最も大きい \(k\) が複数存在する場合には、その \(DP[i][j][k]\) の値とそれを実現する \(k\) の値を \(2\) つ、\(DP[i][j][k]\) の値が \(2\) 番目に大きい \(k\) が複数存在する場合には、最も大きい \(k\) と その \(DP[i][j][k]\) の値のペアと、\(2\) 番目に大きい \(DP[i][j][k]\) の値とそれを実現する \(k\) の値のうちの \(1\) つを記録することにします。

具体的には、\(0\leq i\leq i_0-1\), \(0\leq j\leq k\) について、\(DP[i][j][k]\) \((0\leq k\leq N)\) の値のうち大きい方から\(1,2\) 番目が \(k=k_1(i,j)\) のときと \(k=k_2(i,j)\) のときだったとすると、これらの値から次のように\(DP[i_0][j][k]\) \((0\leq k\leq N)\) の値のうち大きい方から\(1,2\) 番目を求める事ができます。

\(j=0\) のときは \(k\neq C_{i_0}\) ならば \(-\infty\) となるため、最も大きい値は \(C_{i_0-1}\neq C_{i_0}\) ならば \(k=C_{i_0}\) のときに \((i,j)=(i_0-1,0)\)の時の\(dp[i][j][k]\) の最大値に \(V_{i_0}\) を加えた値(\(dp[i][j][k]=-\infty\) ならば \(-\infty\))をとり、\(2\) 番目に大きい値は \(-\infty\) (と\(C_{i_0}\) 以外の適当な \(k\) のペア)となります。
そうでないとき、\(k\neq C_{i_0},k_1(i_0-1,j-1), k_2(i_0-1,j-1)\) であるような \(k\) についての \(DP[i_0][j][k]\) の値は \(DP[i_0][j][k]=DP[i_0-1][j-1][k]\leq DP[i_0-1][j-1][k_2(i_0-1,j-1)]\) となります。 \(C_{i_0}\)\(k_1(i_0-1,j-1)\) または \( k_2(i_0-1,j-1)\) のどちらかと一致するような場合においても、 \(DP[i_0][j][k_1(i_0-1,j-1)]\geq DP[i_0-1][j-1][k_1(i_0-1,j-1)]\geq DP[i_0-1][j-1][k_2(i_0-1,j-1)]\), \(DP[i_0][j][k_2(i_0-1,j-1)]\geq DP[i_0-1][j-1][k_2(i_0-1,j-1)]\) が成り立つため、\(DP[i_0][j][k]\) \((0\leq k\leq N)\) の値のうち大きい方から\(1,2\) 番目となり得るのは \(k=C_{i_0},k_1(i_0-1,j-1), k_2(i_0-1,j-1)\) のいずれかに限られます。
なお、\(DP[i_0][j][C_{i_0}]\) の値は、 \(C_{i_0}\neq k_1(i_0-1,j)\) ならば\(\max(dp[i_0-1][j-1][C_{i_0}],dp[i_0-1][j][k_1(i_0-1,j)]+V_{i_0})\)\(C_{i_0}=k_1(i_0-1,j)\) ならば\(\max(dp[i_0-1][j-1][C_{i_0}],dp[i_0-1][j][k_2(i_0-1,j)]+V_{i_0})\) となります。
このようにして、\(2\) つまたは \(3\) つの\(k\) の候補について \(DP[i_0][j][k]\) の値を比較して大きい方から \(1,2\) 番目を記録すれば良いです。

この時、計算量は各 \((i,j)\) について \(O(1)\) であることから、全体で \(O(NK)\) であり、よって、十分高速にこの問題を解くことができました。

c++ による実装例:

#include <bits/stdc++.h>

using namespace std;

#define rep(i, n) for(int i = 0; i < n; ++i)

#define MAX_K 500
#define INF (long long)2e+18

int main(void){
	int n,k;
	long long c,v;
	long long dp[MAX_K+1][2][2];

	cin>>n>>k;
	rep(j,k+1)rep(ii,2){dp[j][ii][0]=-1;dp[j][ii][1]=-INF;} 
	dp[0][0][0]=0,dp[0][0][1]=0;

	rep(i,n){
		cin>>c>>v;

		for(int j=k;j>=1;j--){
			if(dp[j][0][0]!=c){
				dp[j][0][0]=c;
				dp[j][0][1]+=v;
			}
			else{
				dp[j][0][0]=c;
				dp[j][0][1]=dp[j][1][1]+v;
			}
			dp[j][1][0]=-1;
			dp[j][1][1]=-INF;
			rep(kk,2){
				if(dp[j-1][kk][1]>=dp[j][0][1]){
					if(dp[j-1][kk][0]!=dp[j][0][0]){
						dp[j][1][0]=dp[j][0][0];
						dp[j][1][1]=dp[j][0][1];
					}
					dp[j][0][0]=dp[j-1][kk][0];
					dp[j][0][1]=dp[j-1][kk][1];
				}
				else if((dp[j-1][kk][1]>=dp[j][1][1])&&(dp[j-1][kk][0]!=dp[j][0][0])){
					dp[j][1][0]=dp[j-1][kk][0];
					dp[j][1][1]=dp[j-1][kk][1];
				}
			}
		}
		if(dp[0][0][0]!=c){
			dp[0][0][0]=c;
			dp[0][0][1]+=v;
		}
		else{
			dp[0][0][1]=-INF;
		}
	}

	if(dp[k][0][1]<0)cout<<"-1"<<endl;
	else cout<<dp[k][0][1]<<endl;
	return 0;
}

posted:
last update: