SRM551 Div2 Hard(950) ColorfulCupcakesDivTwo

ColorfulCupcakesDivTwo

動的計画法。現在の位置、最初の人のケーキの色、前の人のケーキの色、各色のケーキの残数、ごとに配り方の数を覚えておく。

#include <string>
#include <vector>
#include <algorithm>
using namespace std;

long long M = 1000000007;
vector<vector<vector<int> > > memo[50][3][3];

//  ケーキ数, 位置, 最初のケーキの色, 1個前のケーキの色, 各色の残数
int BT( int n, int p, int first, int prev, int a, int b, int c )
{
    if ( memo[p][first][prev][a][b][c]>=0 )
        return memo[p][first][prev][a][b][c];

    long long ans = 0;

    if ( p==0 )
    {
        if (a>0)  ans += BT(n,p+1,0,0,a-1,b,c);
        if (b>0)  ans += BT(n,p+1,1,1,a,b-1,c);
        if (c>0)  ans += BT(n,p+1,2,2,a,b,c-1);
    }
    else if ( p<n-1 )
    {
        if (a>0 && prev!=0)  ans += BT(n,p+1,first,0,a-1,b,c);
        if (b>0 && prev!=1)  ans += BT(n,p+1,first,1,a,b-1,c);
        if (c>0 && prev!=2)  ans += BT(n,p+1,first,2,a,b,c-1);
    }
    else
    {
        if (a>0 && prev!=0 && first!=0) ans++;
        if (b>0 && prev!=1 && first!=1) ans++;
        if (c>0 && prev!=2 && first!=2) ans++;
    }
    
    return memo[p][first][prev][a][b][c] = int(ans%M);
}

class ColorfulCupcakesDivTwo{public:
int countArrangements( string cupcakes )
{
    int a = (int)count(cupcakes.begin(),cupcakes.end(),'A');
    int b = (int)count(cupcakes.begin(),cupcakes.end(),'B');
    int c = (int)count(cupcakes.begin(),cupcakes.end(),'C');
    int n = a+b+c;

    for ( int i=0; i<n; i++ )
    for ( int j=0; j<3; j++ )
    for ( int k=0; k<3; k++ )
        memo[i][j][k] = vector<vector<vector<int> > >( a+1,
                        vector<vector<int> >( b+1,
                        vector<int>( c+1, -1 ) ) );
    return BT(n,0,0,0,a,b,c);
}};