SRM564 Div1 Medium "AlternateColors2"
問題
http://community.topcoder.com/stat?c=problem_statement&pm=12345
訳
ボブはボール粉砕ロボットを動かしている。最初、ボブは赤いボールをr個、緑のボールをg個、青のボールをb個持っていた。このロボットはボールがなくなるまで、以下の3ステップを繰り返し行う。
step-1 もし赤のボールがひとつでもあるのなら、赤のボールを1つ壊す。 step-2 もし緑のボールがひとつでもあるのなら、緑のボールを1つ壊す。 step-3 もし青のボールがひとつでもあるのなら、青のボールを1つ壊す。
ボブはそれぞれの色のボールを最初何個持っていたか忘れてしまった。彼は、全部でn個のボールを持っていて、ロボットがk番目に破壊したボールの色は赤であったことのみ覚えている。これらの情報に合うような、最初のボールのセットの場合の数を返せ。正確にはr + g + b = nであり、k番目に壊されるボールが赤であるような、(r, g, b)の互いに異なる数を返せ。
制約
1 <= n <= 10^5
1 <= k <= n
考えた事
dp[ 何個ボールを壊したか(i) ][ RGBのどのボールが残っているか(type) ][ 最後に壊したボールの色(last) ]でDPする。
初期化
ボールを一個も壊しておらず、最初赤のボールを壊したいため、最後に壊したボールの色が青であるようなものが 1
dp[ 0 ][ 任意 ][ B ] = 1
dp[ 0 ][ 任意 ][ R・G ] = 0
更新
次に壊すボールの色をnextとする。
next色のボールを壊したときに、まだnext色のボールが残っている時
dp[ i+1 ][ type ][ next ] += dp[ i ][ type ][ last ]
next色のボールを壊したときに、そこでnext色のボールが全て無くなったとき
残ったボールをnewTypeとすると、
dp[ i+1 ][ newType ][ next ] += dp[ i ][ type ][ last ]
また、i == kであった時は赤が壊されていないといけないため、
lastが赤以外のときは更新はしないようにする。
結果
全てのボールが無くなっていないといけないため、
res = dp[ n ][ NONE ][ R ] + dp[ n ][ NONE ][ G ] + dp[ n ][ NONE ][ B ]
ただし、n == kの時、最後に赤が壊されていないといけないため、
res = dp[ n ][ NONE ][ R ]
ソースコード
long long dp[100050][1<<3][3]; class AlternateColors2 { public: // 次のボールの色を返す int nextColor( int type, int last ){ last = (last + 1) % 3; if( type & (1 << last) ) return last; last = (last + 1) % 3; if( type & (1 << last) ) return last; return (last + 1) % 3; } long long countWays(int n, int k){ // dp初期化 memset( dp, 0, sizeof(dp) ); for( int t = 0; t < 1 << 3; t++ ){ dp[0][t][2] = 1; } // dp更新 for( int i = 0; i < n; i++ ){ for( int type = 0; type < 1 << 3; type++ ){ for( int last = 0; last < 3; last++ ){ if( dp[i][type][last] == 0 ) continue; // k番目が赤でなかった if( i == k && last != 0 ) continue; int next = nextColor( type, last ); // 次のボール // そのボールはもう残ってなかった if( (type & (1 << next)) == 0 ) continue; // nextのボールはまだ残る dp[i+1][type][next] += dp[i][type][last]; // nextのボールを使い切った int newType = type & ~( 1 << next ); dp[i+1][newType][next] += dp[i][type][last]; } } } // 結果 return dp[n][0][0] + (n != k ? (dp[n][0][1] + dp[n][0][2]) : 0); } };
touristさんのやり方
壊されたk個のうち、何個Rが壊されたかを決めて計算していく。
壊されたRをred個とすると、壊されたGB(notred)は k - red個となる
k個を壊すまでのロボットの3ステップぐるぐる数(step)はred - 1回となる
red, notred, stepを使い、以下の3パターンにわけて考える。
1. kまで壊した時に、RGB全て残っている可能性がある場合
step * 2 == notredであれば、ロボットはk個破壊するまでに、RGBRGBRGBRGB...と破壊したはずである。すなわち、k個壊した時に、RGB全て残っている可能性がある。k個まで破壊すると、残りのボールはn - k個となり、これはR,G,Bに自由に振り分けられる。例えるならば、n - k個の区別のつかないお菓子を3人の子供に配る方法の数ということで、これは重複組み合わせになる。
よって、3_N_(n-k) = 3+(n-k)-1_C_3-1 = n-k+2_C_2 = (n-k+2) * (n-k+1) / 2となる
2. kまで壊した時に、GBが残っている可能性がない場合
step > notredであれば、k個壊した時には絶対にRしか残っていない事になる。そのためk~nまでのパターン数は考えなくてよい。GBの振り分け方は次のようになり、それはnotred + 1である。
例) step = 4, notred = 3 R G B ○ ○ ○ ○ ○ ○ ○ ● R G B ○ ○ ○ ○ ○ ○ ○ ● R G B ○ ○ ○ ○ ○ ○ ○ ● R G B ○ ○ ○ ○ ○ ○ ○ ●
3. kまで壊した時に、GBのどちらかが残っている可能性がでてくる場合
step <= notredの時のkまでの振り分け方は次のようになる。
例) step = 4, notred = 5 R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ●
とりあえず、kまでのGBへの振り分け方は、step - (notred - step) + 1となる。
ここで、Rと同じ数だけ壊されるパターンがG,Bで1つずつある。
この場合、k以降、RとG(またはB)に、n - k個のボールを振り分けるられる
例) n - k = 3 R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ (これは数えあげない!!) R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ R G B ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○
上記のように、n - k + 1の振り分け方があるが、BGに一個も振り分けない状態は既に数え上げているため、BGそれぞれでn - k個のパターンがある。
事後
しゅ・・しゅごい・・・
ソースコード
long long countWays(int n, int k){ long long res = 0; for( int red = 1; red <= k; red++ ){ int notred = k - red; int step = red - 1; if( step * 2 < notred ) continue; if( step * 2 == notred ){ res += (long long)(n - k + 2) * (n - k + 1) / 2; }else if( step > notred ){ res += notred + 1; }else{ res += step - (notred - step) + 1; res += (n - k) * 2; } } return res; }