Official

E - Colorful Subsequence Editorial by en_translator


For convenience, we assume that there is a ball with color \(0\) and value \(0\) in the left of the original sequence, and this ball is never removed.
(We call it the \(0\)-th ball from the left.)

First, we may come up with the following simple Dynamic Programming (DP).
For \(0\leq i\leq N\), \(0\leq j\leq K\), and \(0\leq k\leq N\), let \(DP[i][j][k]\) be the maximum total value of remaining balls when considering the \(0\)-th through \(i\)-th balls from the left, removing \(j\) balls from them, the remaining balls have no adjacent pair of balls with the same color, and the rightmost ball has color \(k\). If there is no such way to remove the balls (including when \(i<j\)), let it \(-\infty\).
The final answer is \(\displaystyle \max_{0\leq k\leq N} dp[N][K][k]\).

First, \(DP[0][0][0]=0\), and for \((j,k)\neq (0,0)\), we have \(DP[0][j][k]=-\infty\). Next, we consider the value \(DP[i][j][k]\) for some \(i=i_0\). If \(k\neq C_{i_0}\), then we must remove the \(i_0\)-th ball, so \(DP[i_0][j][k]=DP[i_0-1][j-1][k]\) (or \(-\infty\) if \(j=0\)). If \(k=C_{i_0}\), the maximum value when the \(i_0\)-th ball is removed is \(DP[i_0-1][j-1][k]\) (or \(-\infty\) if \(j=0\)) just as before, and \(\displaystyle \max_{k’\neq k} DP [i_0-1][j][k’]+V_{i_0}\) if the \(i_0\)-th ball is not removed, so \(DP[i][j][k]\) is whichever the larger.
However, implementing this naively costs \(O(N^3K)\) or \(O(N^2K)\) time, which does not fit in the time limit.

Here, carefully inspecting the transition, we realize that it is sufficient to record for each \((i,j)\) only two pairs \((k, [i][j][k])\) with the largest \(DP[i][j][k]\).
Here, if there are multiple \(k\)’s with the largest \(DP[i][j][k]\), we record that \(DP[i][j][k]\) along with two \(k\)’s that achieve it; if there are multiple \(k\)’s with the second largest \(DP[i][j][k]\), we record the largest \(DP[i][j][k]\) along with the \(k\) that achieves it, and the second largest \(DP[i][j][k]\) along with one of the \(k\)’s that achieve that second largest value.

Specifically, for \(0\leq i\leq i_0-1\) and \(0\leq j\leq k\), if the largest and second largest values of \(DP[i][j][k]\) \((0\leq k\leq N)\) are achieved when \(k=k_1(i,j)\) and \(k=k_2(i,j)\), respectively, then the first and second largest values of \(DP[i_0][j][k]\) \((0\leq k\leq N)\) can be found as follows.

When \(j=0\), it takes on \(-\infty\) for \(k\neq C_{i_0}\), so the largest value is the maximum \(dp[i][j][k]\) at \(k=C_{i_0}\), and is obtained as \(V_{i_0}\) plus the maximum \(dp[i][j][k]\) for \((i,j)=(i_0-1,0)\) (or \(-\infty\) if \(dp[i][j][k]=-\infty\)), and the second largest value is \(-\infty\) (paired with some arbitrary \(k\) other than \(C_{i_0}\)).
Otherwise, for \(k\) such that \(k\neq C_{i_0},k_1(i_0-1,j-1),k_2(i_0-1,j-1)\), the value \(DP[i_0][j][k]\) takes on \(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)]\). For \(k\) such that \(k\) coincides with either \(k_1(i_0-1,j-1)\) or \( k_2(i_0-1,j-1)\), it still holds that \(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)]\), so \(DP[i_0][j][k]\) \((0\leq k\leq N)\) is the largest or second largest for the following candidates: \(k=C_{i_0},k_1(i_0-1,j-1), k_2(i_0-1,j-1)\).
Note that \(DP[i_0][j][C_{i_0}]\) equals \(\max(dp[i_0-1][j-1][C_{i_0}],dp[i_0-1][j][k_1(i_0-1,j)]+V_{i_0})\) if \(C_{i_0}\neq k_1(i_0-1,j)\), and \(\max(dp[i_0-1][j-1][C_{i_0}],dp[i_0-1][j][k_2(i_0-1,j)]+V_{i_0})\) if \(C_{i_0}\neq k_1(i_0-1,j)\).
This way, we can compare \(DP[i_0][j][k]\) for two or three candidates of \(k\), and record the largest and second largest among them.

Here, the complexity is \(O(1)\) for each \((i,j)\), for a total of \(O(NK)\), and hence the problem has been solved fast enough.

Sample code in 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: