Codeforces Round #230 (Div. 1) C: Yet Another Number Sequence

codeforces.com

公式解説と異なる方法で解きました。

問題概要

数列 $\{ F_n \}$ を漸化式 $F _ 1 = 1, F _ 2 = 2, F _ i = F _ {i-1} + F _ {i-2} (i \geq 2)$ で定めます。

整数 $N, K$ が与えられるので、 $\displaystyle{ \sum_{i=1}^N F _ i \ i ^ K}$ を $1000000007$ で割った余りを出力してください。

(元の問題文から文字を大文字にしています)

制約

$1 \leq N \leq 10 ^ {17}$
$1 \leq K \leq 40$

解法

制約から、 $K \times K$ などの行列についての行列累乗に持ち込むことを考えます。

まず、 $F _ N \ N ^ K$ を言い換えます。

$\{ F _ n \}$ はフィボナッチ数列なので、 $F _ N$ は $N-1$ 個の球が一列に並んでいて隣接する 2 個の球がともに選ばれないそのうち何個かを選ぶときの選び方の個数に等しくなります。

また、 $N ^ K$ は $N$ 個の球が一列に並んでいて、そのうち一つを選んで整数 $i$ を書くことを $i = 1, 2, \cdots, K$ について行う (同じ球に整数を二回以上書いてもよい) ときの整数の書き方の個数と等しくなります。

球の個数がずれているので、整数のほうについてはどの球にも書かない整数があってもよい (それらは球 $N$ に書いたことにする) とすると、球の個数が揃うので、選び方と書き方の組の個数が $F _ N \ N ^ K$ となります。

ここで、この個数を以下のような DP で求めることを考えます。

$dp[i][j][k] =$ (左から $i$ 個の球を見て、直前の球が ( $j=0$ : 選ばれていなくて $j=1$ : 選ばれていて)、 $1, 2, \cdots, K$ のうち $k$ 個の数を書いたような選び方と書き方の個数)

このとき、DP の遷移は以下のようになります。

$\displaystyle{dp[i+1][0][k2] += dp[i][0][k1] \cdot \binom{K-k1}{k2-k1}}$
$\displaystyle{dp[i+1][0][k2] += dp[i][1][k1] \cdot \binom{K-k1}{k2-k1}}$
$\displaystyle{dp[i+1][1][k2] += dp[i][0][k1] \cdot \binom{K-k1}{k2-k1}}$

この遷移は $i$ によらないので、遷移を $(2K+2) \times (2K+2)$ 行列に変換しそれを $N-1$ 乗することで $F _ N \ N ^ K$ を得ることができます。

この行列を $A$ とおくと、 $\displaystyle{ \sum _ {i=1}^N F _ i \ i ^ K}$ を求めるには $\displaystyle{ \sum _ {i=0}^{N-1} A ^ i}$ を求めればよく、これは行列累乗和を使って求められます。

ですが、これをそのまま実装すると TLE するので、定数倍高速化を行います。

DP 遷移において $k$ は減少することはないので、行列乗算をするときその部分について計算を省略すると、乗算の回数を約 $\displaystyle{\frac{1}{6}}$ にすることができ、AC を得ることができます。

実装

codeforces.com

#include <bits/stdc++.h>
using namespace std;
const long long MOD = 1000000007;
vector<vector<long long>> matmul(vector<vector<long long>> A, vector<vector<long long>> B){
    int N = A.size();
    int K = N / 4;
    vector<vector<long long>> ans(N, vector<long long>(N, 0));
    for (int i1 = 0; i1 < K; i1++){
        for (int j1 = i1; j1 < K; j1++){
            for (int k1 = j1; k1 < K; k1++){
                for (int i = i1; i < N; i += K){
                    for (int j = j1; j < N; j += K){
                        for (int k = k1; k < N; k += K){
                            ans[i][k] += A[i][j] * B[j][k];
                            ans[i][k] %= MOD;
                        }
                    }
                }
            }
        }
    }
    return ans;
}
vector<vector<long long>> matexp(vector<vector<long long>> A, long long b){
    int N = A.size();
    vector<vector<long long>> ans(N, vector<long long>(N, 0));
    for (int i = 0; i < N; i++){
        ans[i][i] = 1;
    }
    while (b > 0){
        if (b % 2 == 1){
            ans = matmul(ans, A);
        }
        A = matmul(A, A);
        b /= 2;
    }
    return ans;
}
int main(){
  long long n;
  int k;
  cin >> n >> k;
  vector<vector<long long>> binom(k + 1, vector<long long>(k + 1, 1));
  for (int i = 2; i <= k; i++){
    for (int j = 1; j < i; j++){
      binom[i][j] = (binom[i - 1][j - 1] + binom[i - 1][j]) % MOD;
    }
  }
  int k1 = k + 1;
  vector<vector<long long>> A(k1 * 4, vector<long long>(k1 * 4, 0));
  for (int i = 0; i <= k; i++){
    for (int j = i; j <= k; j++){
      long long tmp = binom[k - i][j - i];
      A[i][j] = tmp;
      A[i + k1][j] = tmp;
      A[i][j + k1] = tmp;
    }
  }
  for (int i = 0; i < k1 * 2; i++){
    A[k1 * 2 + i][i] = 1;
    A[k1 * 2 + i][k1 * 2 + i] = 1;
  }
  A = matexp(A, n);
  long long ans = 0;
  for (int i = 0; i < k1 * 2; i++){
    ans += A[k1 * 2][i];
  }
  ans %= MOD;
  cout << ans << endl;
}