SRM 403: TheLuckySequence

(問題概要)
4, 7だけで構成された数値をLucky numberとする。
以下を満たす長さlengthの数列は何個あるか

1. ∀i, A[i] is Lucky number
2. ∀i, ∃j, A[i] = numbers[j]
3. ∀i, last digit of A[i] == first digit of A[i + 1]

(解法)

1. numbersのうちのLucky numberだけ残す
2. 初期状態はnumbersのうちどれでもよく、lengthが増えるごとに状態遷移する感じで考える
3. lengthが小さければDPで良いが大きいので行列の累乗を使う

行列の累乗とかいい加減にライブラリ化するべき。

#include <cstdio>
#include <algorithm>
typedef long long ll;

#define REP(i,n) for(int i=0; i<(int)(n); i++)

#include <set>
#include <queue>
#include <iostream>
#include <sstream>

using namespace std;

class TheLuckySequence {
public:
  int count(vector <int> numbers, int length);
};

bool isLucky(int n){
  if(n == 0) return false;
  while(n){
    int d = n % 10;
    if(d != 4 && d != 7) return false;
    n /= 10;
  }
  return true;
}

string toS(int n){ stringstream ss; ss << n; return ss.str(); }

typedef vector<ll> vec;
typedef vector<vec> mat;

const ll mod = 1234567891;

mat operator * (const mat &lhs, const mat &rhs){
  const int n = lhs.size();
  mat ret(n, vec(n));
  REP(i,n) REP(j,n) REP(k,n){
    ret[i][j] += (lhs[i][k] * rhs[k][j]) % mod;
    ret[i][j] %= mod;
  }
  return ret;
}

mat pow(mat m, int k){
  const int n = m.size();
  mat ret(n, vec(n));
  REP(i,n) ret[i][i] = 1;

  while(k > 0){
    if(k & 1) ret = ret * m;
    m = m * m;
    k >>= 1;
  }

  return ret;
}

int TheLuckySequence::count(vector <int> numbers, int length){
  vector<pair<char, char> > luckys;

  sort(numbers.begin(), numbers.end());
  numbers.erase(unique(numbers.begin(), numbers.end()), numbers.end());

  REP(i,numbers.size()) if(isLucky(numbers[i])){
    const string s = toS(numbers[i]);
    luckys.push_back(make_pair(s[0], s[s.size() - 1]));
  }

  const int n = luckys.size();
  mat m(n, vec(n));

  REP(i,n) REP(j,n) if(luckys[i].second == luckys[j].first)
    m[i][j] = 1;

  m = pow(m, length - 1);

  ll ans = 0;
  REP(i,n) REP(j,n){
    ans += m[i][j];
    ans %= mod;
  }

  return (int)ans;
}