RUPC2018 day3 D： 素因数分解の多様性 (The Diversity of Prime Factorization)

問題ページ
Aizu Online Judge Arena

解法

dp[i][j] = (i番目までで今までに使った最大の素数がj) みたいなDPを書きたくなる。しかしこれでは状態数がO(N*(max(q)以下の素数の数))となり間に合わない。ここで連続した2数が指数の方になることはありえないことを考えると、i番目を見ているときに最大の素数の候補はq[i-1]かq[i-2]しか存在しないことがわかる。つまりDPの状態を dp[n][2] と持つことができ、DPの部分がO(N)で解けるようになる。
素数を求めるのはエラトステネスの篩を使ってO(max(q)loglog max(q))で解けるので合計O(max(q)loglog max(q) + N)となる。

ソースコード

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
#define int ll
using VI = vector<int>;
using VVI = vector<VI>;
using PII = pair<int, int>;

#define FOR(i, a, n) for (ll i = (ll)a; i < (ll)n; ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(x) x.begin(), x.end()
#define PB push_back

const ll LLINF = (1LL<<60);
const int INF = (1LL<<30);
const int MOD = 1000000007;

template <typename T> T &chmin(T &a, const T &b) { return a = min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = max(a, b); }
template <typename T> bool IN(T a, T b, T x) { return a<=x&&x<b; }
template<typename T> T ceil(T a, T b) { return a/b + !!(a%b); }
template<class S,class T>
ostream &operator <<(ostream& out,const pair<S,T>& a){
  out<<'('<<a.first<<','<<a.second<<')';
  return out;
}
template<class T>
ostream &operator <<(ostream& out,const vector<T>& a){
  out<<'[';
  REP(i, a.size()) {out<<a[i];if(i!=a.size()-1)out<<',';}
  out<<']';
  return out;
}

int dx[] = {0, 1, 0, -1}, dy[] = {1, 0, -1, 0};

template<unsigned MOD>
class ModInt {
public:
  unsigned x;
  ModInt(): x(0) { }
  ModInt(signed y) : x(y >= 0 ? y % MOD : MOD - (-y) % MOD) {}
  unsigned get() const { return x; }

  // 逆数
  ModInt inv() const {
    ll a = 1, p = x, e = MOD-2;
    while(e > 0) {
      if(e%2 == 0) {p = (p*p) % MOD; e /= 2;}
      else {a = (a*p) % MOD; e--;}
    }
    a %= MOD;
    return ModInt(a);
  }
  // e乗
  ModInt pow(ll e) {
    ll a = 1, p = x;
    while(e > 0) {
      if(e%2 == 0) {p = (p*p) % MOD; e /= 2;}
      else {a = (a*p) % MOD; e--;}
    }
    a %= MOD;
    return ModInt(a);
  }
  // 2のx乗
  ModInt pow2() {
    ll a = 1, p = 2, e = x;
    while(e > 0) {
      if(e%2 == 0) {p = (p*p) % MOD; e /= 2;}
      else {a = (a*p) % MOD; e--;}
    }
    a %= MOD;
    return ModInt(a);
  }

  // Comparators
  bool operator <(ModInt b) { return x < b.x; }
  bool operator >(ModInt b) { return x > b.x; }
  bool operator<=(ModInt b) { return x <= b.x; }
  bool operator>=(ModInt b) { return x >= b.x; }
  bool operator!=(ModInt b) { return x != b.x; }
  bool operator==(ModInt b) { return x == b.x; }

  // increment, decrement
  ModInt operator++() { x++; return *this; }
  ModInt operator--() { x--; return *this; }

  // Basic Operations
  ModInt &operator+=(ModInt that) {
    x = ((ll)x+that.x)%MOD;
    return *this;
  }
  ModInt &operator-=(ModInt that) {
    x = ((((ll)x-that.x)%MOD)+MOD)%MOD;
    return *this;
  }
  ModInt &operator*=(ModInt that) {
    x = (ll)x * that.x % MOD;
    return *this;
  }
  // O(log(mod))かかるので注意
  ModInt &operator/=(ModInt that) {
    x = (ll)x * that.inv() % MOD;
    return *this;
  }
  ModInt &operator%=(ModInt that) {
    x = (ll)x % that.x;
    return *this;
  }
  ModInt operator+(ModInt that)const{return ModInt(*this) += that;}
  ModInt operator-(ModInt that)const{return ModInt(*this) -= that;}
  ModInt operator*(ModInt that)const{return ModInt(*this) *= that;}
  ModInt operator/(ModInt that)const{return ModInt(*this) /= that;}
  ModInt operator%(ModInt that)const{return ModInt(*this) %= that;}
};
typedef ModInt<1000000007> mint;
// Input/Output
ostream &operator<<(ostream& os, mint a) { return os << a.x; }
istream &operator>>(istream& is, mint &a) { return is >> a.x; }

bool prime[1000010];
mint dp[100010][2];
signed main(void)
{
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n; cin >> n;
  VI q(n);
  REP(i, n) cin >> q[i];

  memset(prime, true, sizeof(prime));
  prime[0] = prime[1] = false;
  for (int i = 2; i * i <= 1000010; i++) {
    if (prime[i]) {
      for (int j = 2 * i; j <= 1000010; j += i) {
        prime[j] = false;
      }
    }
  }

  dp[0][1] = (prime[q[0]] ? 1 : 0);
  FOR(i, 1, n) {
    if(prime[q[i-1]]) dp[i][0] += dp[i-1][1];
    if(prime[q[i]] && prime[q[i-1]] && q[i-1] < q[i]) dp[i][1] += dp[i-1][1];
    if(i >= 2 && prime[q[i]] && prime[q[i-2]] && q[i-2] < q[i]) dp[i][1] += dp[i-1][0];
  }
  cout << dp[n-1][0] + dp[n-1][1] << endl;

  return 0;
}