[CF1117D] Magic Gems

[CF1117D] Magic Gems - 矩阵乘法优化dp

Description

一颗魔法宝石可以分解为 m 颗普通宝石，魔法宝石和普通宝石都占据 1 体积的空间。现在有一大堆带编号的魔法宝石，可以选出一部分宝石，然后指定一种一部分魔法宝石分解，求有多少种选出体积为 n 的集合的方案。(n le 10^{18}, m le 100)

Solution

设 (f[i]) 表示用了 (i) 个单位空间时的方案数，考虑最后一个选出的魔法宝石是否分解，于是有 (f[i]=f[i-1]+f[i-m])，由于 (m) 很小，矩阵快速幂即可

#include <bits/stdc++.h>
using namespace std;

#define int long long

const int mod = 1e9 + 7;

struct Matrix
{
    int n, m;
    vector<vector<int>> a;

    Matrix(int n, int m) : n(n), m(m)
    {
        a.resize(n);
        for (int i = 0; i < n; i++)
            a[i].resize(m);
    }

    friend Matrix operator*(Matrix lhs, Matrix rhs)
    {
        assert(lhs.m == rhs.n);
        Matrix res(lhs.n, rhs.m);
        for (int i = 0; i < res.n; i++)
        {
            for (int j = 0; j < res.m; j++)
            {
                for (int k = 0; k < lhs.m; k++)
                {
                    res.a[i][j] += lhs.a[i][k] * rhs.a[k][j];
                    res.a[i][j] %= mod;
                }
            }
        }
        return res;
    }

    Matrix operator*=(const Matrix &rhs)
    {
        return (*this) = (*this) * rhs;
    }

    vector<int> &operator[](int x) { return a[x]; }
};

signed main()
{
    int n, m;
    cin >> n >> m;

    Matrix mat(m, m);
    mat[0][0] = mat[0][m - 1] = 1;
    for (int i = 1; i < m; i++)
        mat[i][i - 1] = 1;

    Matrix vec(m, 1);
    vec[0][0] = 1;

    Matrix trans(m, m);
    for (int i = 0; i < m; i++)
        trans[i][i] = 1;
    while (n)
    {
        if (n & 1)
            trans *= mat;
        mat *= mat;
        n /= 2;
    }

    vec = trans * vec;

    cout << vec[0][0] << endl;
}