kmjp's blog

競技プログラミング参加記です

AtCoder ARC #020 : D - お菓子の国の旅行

これは本番には思い浮かばないな…。本番は部分点のみ。
http://arc020.contest.atcoder.jp/tasks/arc020_4

問題

N個の町が数直線上で左から順に並んでおり、それぞれの位置が与えられる。

N個中K個の町を以下のように訪問したとき、距離の合計がMの倍数となる組み合わせを答えよ。

  • 訪問は任意の町で初めて任意の町で終わることができる。
  • 訪問する町の間は最短距離で移動しなければならない。
  • 訪問順が違う訪問方法は別物とカウントする。
  • K個の町は訪問は1回ずつしかできないが、通過回数は何回でも良い。

解法

部分点ではN<=12、K<=10と数が少ないので、通った町のBitDPやDFSでも何とか間に合う。
自分もDFS+半分全列挙で通した。

満点はN<=100のためそのままBitDPやDFSでは間に合わない。
そこで以下の通り解説を見てチャレンジ。
AtCoder Regular Contest 020 解説

N個の町の間に(N-1)個の道があるので、これを端からDPしていく。

各道に対する状態として、K個の訪問する町が道の左右どちらにあるかをbitmaskで持とう。
各道を通る回数は、その道の左から右または右から左の町に移動する回数なので、bitmask中隣接するビットが異なっている数だけ移動する。

i番目の道の状態に対応する組み合わせ数をdp[道番号][bitmask][距離の総和%M]とする。
i番の道のDP結果に対して、i+1番の道では以下の手順が考えられる。

  • i番と同じbitmaskを保持する。これは2つの道の間にある町(i+1番の町)を訪問するケースがない場合を示す。
  • i番目のbitmaskのうちj bit目が「右の町」となっている場合、そこを「左の町」と直す。これはj番目に2つの道の間にある町(i+1番の町)で訪問するケースを示す。

それぞれの場合について、道を通る回数を用いて[距離の総和%M]を更新していけばよい。
なお、下記のコードでは最後に(N-1)番目と存在しないN番目の町をつなぐ存在しない距離0でN番目の道を番兵としておいている。
どのような移動経路にしろ、この道に対してはすべての訪問した町が左側になければならないので、答えとしてはdp[N][全ビットが左側の町を示すbitmask][0]となる。

int N,M,K;
int D[101];
ll mo=1000000007;
ll dp[101][1024][31];
ll ret=0;
int dm[1024];


void solve() {
	int f,i,j,k,l,x,y,mask,mask2;
	
	cin>>N>>M>>K;
	FOR(i,N-1) cin>>D[i];
	
	
	FOR(i,1024) FOR(x,K-1) dm[i]+=((i>>x)^(i>>(x+1)))&1;
	
	dp[0][0][0]=1;
	FOR(i,N) {
		FOR(mask,1<<K) {
			FOR(x,M) {
				dp[i+1][mask][(x+D[i]*dm[mask])%M] += dp[i][mask][x];
				dp[i+1][mask][(x+D[i]*dm[mask])%M] %= mo;
				FOR(j,K) {
					if((mask&(1<<j))==0) {
						mask2 = mask | (1<<j);
						dp[i+1][mask2][(x+D[i]*dm[mask2])%M] += dp[i][mask][x];
						dp[i+1][mask2][(x+D[i]*dm[mask2])%M] %= mo;
					}
				}
			}
		}
	}
	
	cout << dp[N][(1<<K)-1][0] << endl;
}

まとめ

答えだけ見ると、部分点解法よりよっぽど楽だ。