ARC074 E. RGB Sequence - phyllo’s algorithm note

問題

N個のマスが横一列に並んでおり、各マスを赤、緑、青のいずれかで塗りたい。
しかし、 M個の「マスlからマスrの色の種類数がx」という条件を満たすようにしたい。
条件を満たすような塗り方は何通りか？10^9+7の余りを答えよ。

制約

1 <= N <= 300
1 <= M <= 300
1 <= l_i <= r_i <= N
1 <= x_i <= 3

解説

Editorial - AtCoder Regular Contest 074 | AtCoder
AtCoder Regular Contest 074/ Beginner Contest 062 解説放送 - YouTube
以下、だいぶ自分の解釈が入っているので、間違っているかもしれないことに注意。

種類数を数えたいので、DPが適用できないか考える。

条件がない場合の数え上げ

まず、M個の条件がない場合を考える。
左から順番に塗っていって、i番目までの計算結果があるとする。
このとき、
dp[i][...]:=i番目まで塗ったときの塗り方の総数
を考える。
このとき「赤を最後に塗った場所がr」「緑を最後に塗った場所がg」「青を最後に塗った場所がb」というのがわかっていれば、
i+1に赤を使った場合は、dp[i+1][i+1][g][b] += dp[i][r][g][b]、
i+1に緑を使った場合は、dp[i+1][r][i+1][b] += dp[i][r][g][b]、
i+1に青を使った場合は、dp[i+1][r][g][i+1] += dp[i][r][g][b]、
と更新することができる。
しかし、O(N^4)になってしまっている。が、これはだいぶ無駄で、i番目のとき、rまたはgまたはbのどれか一つはiになっているところだけ考えればよいはずなので、iは必要ない。

dp[r][g][b]のとき、i=max(r,g,b)と計算できるので、
i+1に赤を使った場合は、dp[i+1][g][b] += dp[r][g][b]
i+1に緑を使った場合は、dp[r][i+1][b] += dp[r][g][b]
i+1に青を使った場合は、dp[r][g][i+1] += dp[r][g][b]
と更新できる。
これはO(N^3)なので、間に合う。

条件を満たさないとき、数え上げない

条件を満たすときだけの数え上げをしたいので、dpの更新中にこの判定をいれたい。

(r,g,b)が決まっているとき、i=max(r,g,b)とすると、M個の各条件(l_i,r_i,x_i)について、「l_i<=r<=r_i」「l_i<=g<=r_i」「l_i<=b<=r_i」であるか確認して使われている色の種類数を確認して、x_iと一致するかどうか見ればよい。
そして、条件を満たさない場合はdp[r][g][b]=0としてしまえばよい。

しかし、r,g,bの三重ループの中で単純にM個のループを回してしまうとO(M*N^3)となる。
そこで、「vector < pair < int,int > > v」として、v[r_i][j] = pair(l_i,x_i)のように持たせて、各iでv[i]だけをチェックする。
iは0～Nの範囲で動くが、v[i]のようにM個の条件をiで分割して持たせると、O(N+M)になる。
全体ではO(M*N^3)はO(N^2*(N+M))=O(N^3 + M*N^2)となり、間に合う。

コード

#define MOD 1000000007

int dp[305][305][305];
vector<pair<int,int>> v[305];

int main(){
  int N, M;
  cin >> N >> M;
  rep(i,M){
    int l, r, x;
    cin >> l >> r >> x;
    v[r].push_back(make_pair(l,x));
  }

  ll ret = 0;
  dp[0][0][0] = 1;

  rep(i,N+1){
    rep(j,N+1){
      rep(k,N+1){
        int m = max(i,max(j,k));

        rep(ii,v[m].size()){
          int r = m;
          int l = v[r][ii].first;
          int x = v[r][ii].second;

          int cnt = 0;
          if(l<=i && i<=r) cnt++;
          if(l<=j && j<=r) cnt++;
          if(l<=k && k<=r) cnt++;
          if(cnt != x){
            dp[i][j][k] = 0;
          }
        }

        if(m == N){
          ret += dp[i][j][k];
          ret %= MOD;
          continue;
        }

        dp[m+1][j][k] += dp[i][j][k];
        dp[m+1][j][k] %= MOD;
        dp[i][m+1][k] += dp[i][j][k];
        dp[i][m+1][k] %= MOD;
        dp[i][j][m+1] += dp[i][j][k];
        dp[i][j][m+1] %= MOD;
      }
    }
  }
  
  cout << ret << endl;

  return 0;
}

反省

条件を持つ数え上げとかは、桁DPとかでもよく見る
無駄な状態を減らす
O(N)でM個の条件のチェックをする場合、位置に依存するならばO(NM)ではなくO(N+M)に落とせる