SRM551 Div2 Hard(950) 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); }};