Div2 500 EllysPlaylists

問題

3文字以上の，文字列の集合（文字列全体が同じものは無い）が与えられている．その中から重複を含めて K 個抽出して並べるパターンは(mod 1000000007で)何通りあるか．
ただし，並べた順が s₀, s₁, …, s_K-1 のとき，s_i の末尾３文字は s_i+1 の頭３文字と一致していなくてはならない．

考え方

DPってやつ？をやれば大丈夫．つまり各文字列スタートで k 個並べるパターンの数を数えて，k+1 個のときにどうなる，ということを k=1..K で計算する．

コード

#include <vector>
#include <string>
using namespace std;
typedef long long lint;
const lint MOD = 1000000007LL;

class EllysPlaylists
{
public:
  int countPlaylists(vector <string> songs, int K)
  {
    const int N = songs.size();
    vector< vector<bool> > cont(N, vector<bool>(N, false) );
    for( int i = 0 ; i < N ; ++i ) {
      string end = songs[i].substr(songs[i].size()-3);
      for( int j = 0 ; j < N ; ++j ) {
        if( end == songs[j].substr(0, 3) ) {
          cont[i][j] = true;
        }
      }
    }
    
    vector<lint> cnt(N, 1);
    for( int k = 1 ; k < K ; ++k ) {
      vector<lint> nxt(N, 0);
      for( int i = 0 ; i < N ; ++i ) {
        lint c = 0;
        for( int j = 0 ; j < N ; ++j ) {
          if( cont[i][j] ) {
            c = (c + cnt[j]) % MOD;
          }
        }
        nxt[i] = c;
      }
      cnt = nxt;
    }
    lint sum = 0;
    for( int i = 0 ; i < N ; ++i ) sum = (sum + cnt[i]) % MOD;
    return sum;
  }
};