SRM 227 Div1 Medium TreeSpreading

問題

4つの種類の木があり一列に並べる.同じ種類の木を連続して配置しない時,可能な並べ方の個数を返せ.

解答

メモ化再帰.dp[a][b][c][d]にそれぞれの種類の個数がa,b,c,dの時の並べ方の個数を格納しておく.この時の個数はdp[a-1][b][c][d] + dp[a][b-1][c][d] + dp[a][b][c-1][d] + dp[a][b][c][d-1]となる(a,b,c,d>0の時).もしも一つ前の木の種類がaだった場合はaを,全体からaを選択した時の数を引いてやれば求めることが出来る.dp[a][b][c][d] - dp[a-1][b][c][d].

class TreeSpreading
{
public:
  long long dp[11][11][11][11];
  int inf;
  long long calc(vector<int> & state, int prev){
    // cout << "prev=" << prev << endl;
    vector<int> cur = state;
    int a = state[0];
    int b = state[1];
    int c = state[2];
    int d = state[3];
    if (dp[a][b][c][d] != inf){
      if (prev != inf){
        if (cur[prev] > 0){
          cur[prev]--;
          return dp[a][b][c][d] - calc(cur, prev);
          // return dp[a][b][c][d] - dp[cur[0]][cur[1]][cur[2]][cur[3]];
        }else{
          return dp[a][b][c][d];
        }
      }
      return dp[a][b][c][d];
    }
    if (a == 0 && b == 0 && c==0 && d==0){
      dp[a][b][c][d] = 1;
      return 1;
    }
    long long total = 0;
    for(int i = 0; i < 4; i++){
      if (cur[i] <= 0) continue;
      cur[i]--;
      total += calc(cur, i);
      cur[i]++;
    }
    dp[a][b][c][d] = total;
    if (prev != inf && cur[prev] > 0){
      cur[prev]--;
      total -= calc(cur, prev);
      cur[prev]++;
    }
    return total;
  }
  long long countArrangements(int a, int b, int c, int d)
    {
      inf = -1;
      vector<int> state;
      state.push_back(a); state.push_back(b);
      state.push_back(c); state.push_back(d);
      for(int i1 = 0; i1 < 11; i1++){
      for(int i2 = 0; i2 < 11; i2++){
      for(int i3 = 0; i3 < 11; i3++){
      for(int i4 = 0; i4 < 11; i4++){
        dp[i1][i2][i3][i4] = inf;
      }
      }
      }
      }
      return calc(state, inf);
    }
};