SRM 551 Div2 Hard ColorfulCupcakesDivTwo

問題

http://apps.topcoder.com/stat?c=problem_statement&pm=12138&rd=15173
複数のカップケーキがあり、'A', 'B', 'C'のいずれかの色がついている。これらのカップケーキを円形に並べたい。隣合うカップケーキが同じ色にならないように並べるとき、何通りの並べ方が存在するか。1000000007で割ったあまりを返せ。

やりかた

dp[a][b][c][s][p] := （並べはじめのカップケーキの色がsで、並べ終わりのカップケーキの色がpで、'A'のカップケーキがa残っており、'B'のカップケーキがb残っており、'C'のカップケーキがc残っているような並べ方の総数）でDP。

再帰による探索でa, b, cの許す限り、pとは異なる色のケーキを置いていくケースを探索すればいい。カップケーキをすべて置き終わったときのみ、sとpが異なる必要があることに注意。

antaさんの記事、Editorialを参考にしました。

以下ソース。

const ll MOD = 1000000007;
int num[3];
ll dp[51][51][51][3][3];

class ColorfulCupcakesDivTwo {
public:
  ll rec(int a, int b, int c, int s, int p){
    if(dp[a][b][c][s][p] != -1) return dp[a][b][c][s][p];

    ll ret = 0;
    if(a + b + c == 0){//ケーキをすべて置き終わったとき
      ret = (s != p) ? 1 : 0;
    }else{
      if(a > 0 && p != 0) ret = (ret + rec(a - 1, b, c, s, 0)) % MOD;
      if(b > 0 && p != 1) ret = (ret + rec(a, b - 1, c, s, 1)) % MOD;
      if(c > 0 && p != 2) ret = (ret + rec(a, b, c - 1, s, 2)) % MOD;
    }
    return dp[a][b][c][s][p] = ret % MOD;
  }
  int countArrangements(string cupcakes) {
    memset(dp, -1, sizeof(dp));
    memset(num, 0, sizeof(num));
    for(int i = 0; i < (int)cupcakes.length(); i++)
      num[cupcakes[i] - 'A']++;

    ll ans = 0;
    if(num[0] > 0) ans = (ans + rec(num[0] - 1, num[1], num[2], 0, 0)) % MOD;
    if(num[1] > 0) ans = (ans + rec(num[0], num[1] - 1, num[2], 1, 1)) % MOD;
    if(num[2] > 0) ans = (ans + rec(num[0], num[1], num[2] - 1, 2, 2)) % MOD;
    return ans % MOD;
  }

シンプルにやればできたかもしれない。