AtCoder ABC #440 (Promotion of Engineer Guild Fes) : G - Haunted House

これは解けても良かったな…。
https://atcoder.jp/contests/abc440/tasks/abc440_g

問題

F階建ての建物があり、F*H*Wのグリッドで表現される。
各マスは、移動できないマスかできるマスのいずれかで、移動可能なマスはそれぞれコインが置かれておりその数が示される。
グリッド上のマスの位置を１つ指定されるので、以下の問いに答えよ。

• 指定されたマスから始め、以下の移動を繰り返しながら到達したマスのコインを回収するとする
  • 同フロア内で、隣接する移動可能マスに移動する
  • １つ真上の移動可能マスに移動する
  • １つ真下の移動可能マスに移動する。ただしこれは１回しか行えない

これを繰り返し到達可能なマスにおけるコインの回収枚数の総和の最大値を求めよ。

解法

まず同フロア内で移動可能なマス群は、Union-Findで連結させ、そのコインの総和を求めよう。
以下この連結成分を部屋と呼ぶ。

ここから以下の順で値を求めて行く。

• dp(r) := 部屋r内にあるコインの総和
• dp2(r) := 部屋rから初めて、真下移動を行わない場合の、コインの回収枚数の総和の最大値
• dp3(r) := 部屋rから初めて、この部屋内で真下移動を行うか、または全く真下移動を行わない合の、コインの回収枚数の総和の最大値
• dp4(r) := 部屋rから初めて、今後最大１回まで真下移動を行う場合の、コインの回収枚数の総和の最大値

なお、dp2,dp3,dp4は最大値に上位２位まで覚えておき、かつそれぞれ次に移動する部屋を覚えておこう。
これにより、同じ部屋を２回通過する際、そのコインを二重カウントしないようにできる。

int F,H,W;
string S[10][505];
template<int um> class UF {
	public:
	vector<int> par,rank,cnt,G[um];
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit(int num=um) {int i; FOR(i,num) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<505*505*10> uf;

int dp[10*500*500];
vector<pair<int,int>> dp2[10*500*500];
vector<pair<int,int>> dp3[10*500*500];
vector<pair<int,int>> dp4[10*500*500];
set<int> up[10*500*500],down[10*500*500];

int Q,G,A,B;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>F>>H>>W;
	FOR(k,F) {
		FOR(y,H) {
			cin>>S[k][y];
			FOR(x,W) if(S[k][y][x]!='#') {
				if(y&&S[k][y-1][x]!='#') uf(k*500*500+y*500+x,k*500*500+(y-1)*500+x);
				if(x&&S[k][y][x-1]!='#') uf(k*500*500+y*500+x,k*500*500+y*500+x-1);
			}
		}
	}
	FOR(k,F) {
		FOR(y,H) FOR(x,W) if(S[k][y][x]!='#') {
			dp[uf[k*500*500+y*500+x]]+=S[k][y][x]-'0';
			if(k&&S[k-1][y][x]!='#') {
				up[uf[(k-1)*500*500+y*500+x]].insert(uf[k*500*500+y*500+x]);
				down[uf[k*500*500+y*500+x]].insert(uf[(k-1)*500*500+y*500+x]);
			}
		}
	}
	//梯子無しの場合の最大値
	for(k=F-1;k>=0;k--) FOR(y,H) FOR(x,W) if(S[k][y][x]!='#'&&uf[k*500*500+y*500+x]==k*500*500+y*500+x) {
		int v=k*500*500+y*500+x;
		dp2[v].push_back({dp[v],-1});
		FORR(u,up[v]) dp2[v].push_back({dp[v]+dp2[u][0].first,u});
		sort(ALL(dp2[v]));
		reverse(ALL(dp2[v]));
		if(dp2[v].size()>2) dp2[v].resize(2);
		/*
		cout<<v<<" "<<dp[v]<<" ";
		FORR2(a,b,dp2[v]) cout<<a<<":"<<b<<" ";
		cout<<endl;
		*/
	}
	//今すぐ梯子を使う場合の最大値
	FOR(k,F) FOR(y,H) FOR(x,W) if(S[k][y][x]!='#'&&uf[k*500*500+y*500+x]==k*500*500+y*500+x) {
		int v=k*500*500+y*500+x;
		dp3[v].push_back(dp2[v][0]);
		if(dp2[v].size()==2) dp3[v].push_back(dp2[v][1]);
		FORR(d,down[v]) {
			if(dp2[d][0].second!=v) dp3[v].push_back({dp[v]+dp2[d][0].first,d});
			else {
				//戻ってくる
				//戻らない
				dp3[v].push_back({max(dp2[v][0].first+dp[d],dp[v]+dp2[d][1].first),d});
			}
		}
		sort(ALL(dp3[v]));
		reverse(ALL(dp3[v]));
		if(dp3[v].size()>2) dp3[v].resize(2);
		/*
		cout<<v<<" ";
		FORR2(a,b,dp3[v]) cout<<a<<":"<<b<<" ";
		cout<<endl;
		*/
	}
	//途中で1回梯子を使う場合の最大値
	for(k=F-1;k>=0;k--) FOR(y,H) FOR(x,W) if(S[k][y][x]!='#'&&uf[k*500*500+y*500+x]==k*500*500+y*500+x) {
		int v=k*500*500+y*500+x;
		dp4[v].push_back(dp3[v][0]);
		if(dp3[v].size()==2) dp4[v].push_back(dp3[v][1]);
		FORR(u,up[v]) {
			if(dp4[u][0].second!=v) dp4[v].push_back({dp[v]+dp4[u][0].first,u});
			else {
				//2つ目の候補に行く
				int c1=dp[v]+dp4[u][1].first;
				//戻ってくる
				int c2=dp2[v][dp2[v][0].second==u].first+dp[u];
				dp4[v].push_back({max(c1,c2),u});
			}
		}
		sort(ALL(dp4[v]));
		reverse(ALL(dp4[v]));
		if(dp4[v].size()>2) dp4[v].resize(2);
		/*
		cout<<v<<" ";
		FORR2(a,b,dp4[v]) cout<<a<<":"<<b<<" ";
		cout<<endl;
		*/
	}
	cin>>Q;
	while(Q--) {
		cin>>G>>A>>B;
		G--,A--,B--;
		cout<<dp4[uf[G*500*500+A*500+B]][0].first<<endl;
	}
	
	
}

まとめ

落ち着けば思いつけそうな問題。