yukicoder : No.1515 Making Many Multiples

遠回りな解法をしてしまった。
https://yukicoder.me/problems/no/1515

問題

N+2枚のカードがある。
初期状態でプレイヤーはX,Yの値のある2枚のカードを持ち、あとN枚のカードが山に積んである。
カードに書かれた値はわかっている。

山の一番上のカードを取り、手持ちの3枚のカードの総和がKの倍数なら1ポイントを得る。
その後、3枚中1枚のカードを捨てる。
最適なカードの捨てかたを行うとき、得られるポイントの総和の最大値を求めよ。

解法

カードの値はいずれもmod Kを取った値を考えればよい。

以下の3つの配列を、カードを1枚とるたびに更新していこう。
P(x,y) := 手持ちのカードの値がx,yの時のポイントの最大値
R(x) := P(x,*)の最大値
C(y) := P(*,y)の最大値

山から取ったカードの値がzであるとき、以下の選択肢がある。

• カードを取ってすぐ捨てる。この際、x+y+z=0 (mod K)であるP(x,y)がインクリメントされる。
• 手持ちのyのカードを捨てる。この時、P(x,z)=max(R(x),P(x,y')+1) (x+y'+z=0 mod K)
• 手持ちのxのカードを捨てる。この時、P(z,y)=max(C(y),P(x',y)+1) (x'+y+z=0 mod K)

Pのうち更新されるのはカード枚にO(K)箇所程度なので、この問題はO(NK+K^2)で解ける。

int R[2020],C[2020];


int N,K,X,Y;
int A[2020];
int ma[2020][2020];

void update(int y,int x,int v) {
	if(ma[y][x]<v) {
		ma[y][x]=v;
		R[y]=max(R[y],v);
		C[x]=max(C[x],v);
	}
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	FOR(x,2020) FOR(y,2020) ma[x][y]=-1<<20;
	FOR(x,2020) R[x]=C[x]=-1<<20;
	
	cin>>N>>K>>X>>Y;
	update(X%K,Y%K,0);
	
	
	FOR(i,N) {
		cin>>j;
		j%=K;
		vector<int> cu,ru;
		FOR(x,K) {
			y=(3*K-x-j)%K;
			cu.push_back(max(C[x],ma[x][y]+1));
		}
		FOR(y,K) {
			x=(3*K-j-y)%K;
			ru.push_back(max(R[y],ma[x][y]+1));
		}
		FOR(x,K) {
			y=(3*K-x-j)%K;
			update(x,y,ma[x][y]+1);
		}
		FOR(x,K) update(j,x,cu[x]);
		FOR(y,K) update(y,j,ru[y]);
	}
	int ret=0;
	FOR(x,K) FOR(y,K) ret=max(ret,ma[y][x]);
	cout<<ret<<endl;
	
}

まとめ

本番、最初R,Cを無駄にSegTreeを使いTLEしてしまった。