WARush

SRMの結果とか、解けた問題のコードを書いていきます

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;
}