SRM564 Div1 Medium

500

DP[何個目][どの色が残っている][次に消すべき色] = パターン数
k が 1-based ということを見落として数時間を浪費して悲しい。

const int N = 100000 + 5;
const int BIT = 1 << 3;
const int color = 3;
static lli dp[N][BIT][color];

class AlternateColors2 {
public:
  long long countWays(int n, int k)
  {
    fill(&dp[0][0][0], &dp[N - 1][BIT - 1][color - 1] + 1, 0);
    const int red = 0;
    for (int bit = 1; bit < BIT; ++bit) {
      dp[0][bit][red] = 1;
    }

    for (int nth = 0; nth < n; ++nth) {
      for (int rest = 0; rest < BIT; ++rest) {
        for (int c = 0; c < color; ++c) {
          if (nth == k - 1 && c != red) continue;
          int d = c;
          for (int i = 0; i < 3; ++i) {
            d = (d + 1) % color;
            if (rest & (1 << d)) break;
          }
          int curr = dp[nth][rest][c];
          if (rest & (1 << c)) {
            dp[nth + 1][rest][d] += curr;
            dp[nth + 1][rest ^ (1 << c)][d] += curr;
          }
        }
      }
    }

    return dp[n][0][0] + dp[n][0][1] + dp[n][0][2];
  }
};