Skip to content

Instantly share code, notes, and snippets.

@tosaka2
Created February 1, 2018 17:52
Show Gist options
  • Save tosaka2/e1b4e9dcb7892568d16570075d85941a to your computer and use it in GitHub Desktop.
Save tosaka2/e1b4e9dcb7892568d16570075d85941a to your computer and use it in GitHub Desktop.
kitamasa法でm項間漸化式の第n項までの和を求めるやつ
#include <bits/stdc++.h>
using namespace std;
const int64_t MOD = 1000000007;
// O(mn) m+1項間漸化式のn要素目までの和を前から順に項を計算し足していくことで求める
int64_t simple_sum(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n);
// O(m^3 log(n)) m+1項間漸化式のn要素目までの和を累乗和で求める
int64_t power_series_sum(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n);
// O(m^2 log(n)) m+1項間漸化式のn要素目までの和を総和の漸化式+kitamasa法で求める(今回のメイン)
int64_t kitamasa_sum(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n);
int dummy = 0;
// 時間計測関数
template <typename F>
auto timeit(F f, string func_name, int num_exec, const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n) {
auto start = std::chrono::system_clock::now();
for (int i = 0; i < num_exec; i++) {
dummy = f(coefficients, firsts, n);
}
auto end = std::chrono::system_clock::now();
auto dur = end - start;
auto usec = std::chrono::duration_cast<std::chrono::microseconds>(dur).count();
cout << func_name << ": " << (usec / (double)num_exec / 1000) << " ms" << endl;
}
int main() {
vector<int64_t> firsts(100); // 初項 a_1, a_2, ... , a_m
vector<int64_t> coefficients(100); // 係数 c_1, c_2, ..., c_m
for (int i = 0; i < (int)firsts.size(); i++) {
firsts[i] = i + 1;
coefficients[i] = i + 1;
}
int64_t n = 100000000; // n項目までの和を求める
int num_exec = 10; // 計測回数
//timeit(simple_sum, "simple_sum", num_exec, coefficients, firsts, n);
timeit(power_series_sum, "power_series_sum", num_exec, coefficients, firsts, n);
timeit(kitamasa_sum, "kitamasa_sum", num_exec, coefficients, firsts, n);
}
/*
----------------------------------------------------------
今回のメイン
m+1項間漸化式を総和の漸化式(m+2項間)に変形
*/
// O(m) 総和の漸化式の係数を求める
vector<int64_t> coefficients_sums(const vector<int64_t>& coefficients) {
int m = coefficients.size();
vector<int64_t> new_coefficients(m + 1);
new_coefficients[0] = -coefficients[0];
for (int i = 0; i + 1 < m; i++) {
new_coefficients[i + 1] = coefficients[i] - coefficients[i + 1];
}
new_coefficients[m] = coefficients[m - 1] + 1;
return new_coefficients;
}
// O(m) m+1項間漸化式の初項からm+1項目までを求め、配列で返す
vector<int64_t> first_sums(const vector<int64_t>& coefficients, const vector<int64_t>& firsts) {
int m = firsts.size();
vector<int64_t> sums(m + 1);
sums[0] = firsts[0];
for (int i = 1; i < m; i++) {
sums[i] = (firsts[i] + sums[i - 1]) % MOD;
}
int64_t next = 0;
for (int i = 0; i < m; i++) {
next = (next + (coefficients[i] * firsts[i])) % MOD;
}
sums[m] = (next + sums[m - 1]) % MOD;
return sums;
}
// O(m^2 log(n)) m+1項間漸化式のn要素目をkitamasa法で求める
int64_t kitamasa(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n);
int64_t kitamasa_sum(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n) {
auto a_sums = first_sums(coefficients, firsts);
auto c_sums = coefficients_sums(coefficients);
return kitamasa(c_sums, a_sums, n);
}
/*
----------------------------------------------------------
単純に前から計算
*/
// coefficients : m+1項間漸化式の係数c_1 ~ c_mを与える
// firsts : m+1項間漸化式の初項a_1 ~ a_mを与える
// n : n項目までの和を求める
int64_t simple_sum(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n) {
int m = firsts.size();
deque<int64_t> nexts(m);
int64_t sum = 0;
for (int i = 0; i < m; i++) {
sum = (sum + firsts[i]) % MOD;
nexts[i] = firsts[i];
}
for (int i = m; i < n; i++) {
int64_t next = 0;
for (int j = 0; j < m; j++) {
next = (next + coefficients[j] * nexts[j]) % MOD;
}
sum = (sum + next) % MOD;
nexts.pop_front();
nexts.push_back(next);
}
return sum;
}
/*
----------------------------------------------------------
累乗和
出典:
プログラミングコンテストチャレンジブック第2版
P180 ~ P185
一部改変
*/
using vec = vector<int64_t>;
using mat = vector<vector<int64_t>>;
mat mul(mat &A, mat &B) {
mat C(A.size(), vec(B[0].size()));
for (int i = 0; i < (int)A.size(); i++) {
for (int k = 0; k < (int)B.size(); k++) {
for (int j = 0; j < (int)B[0].size(); j++) {
C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % MOD;
}
}
}
return C;
}
mat pow(mat A, int64_t n) {
mat B(A.size(), vec(A.size()));
for (int i = 0; i < (int)A.size(); i++) {
B[i][i] = 1;
}
while (n > 0) {
if (n & 1) B = mul(B, A);
A = mul(A, A);
n >>= 1;
}
return B;
}
mat make_companion(const vector<int64_t>& coefficients) {
int m = coefficients.size();
mat A(m, vec(m));
for (int i = 0; i < m; i++) {
A[0][i] = coefficients[m - i - 1];
}
for (int i = 0; i < m - 1; i++) {
A[1 + i][i] = 1;
}
return A;
}
mat matrix_power_series(const vector<int64_t>& coefficients, int n) {
int m = coefficients.size();
auto A = make_companion(coefficients);
mat B(m * 2, vec(m * 2));
for (int i = 0; i < m; i++) {
for (int j = 0; j < m; j++) {
B[i][j] = A[i][j];
}
B[m + i][i] = B[m + i][m + i] = 1;
}
B = pow(B, n);
mat S(m, vec(m));
for (int i = 0; i < m; i++) {
for (int j = 0; j < m; j++) {
S[i][j] = B[m + i][j];
}
}
return S;
}
int64_t power_series_sum(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n) {
auto S = matrix_power_series(coefficients, n);
int m = firsts.size();
int64_t sum = 0;
for (int i = 0; i < m; i++) {
sum = (sum + S[m - 1][i] * firsts[m - i - 1]) % MOD;
}
return sum;
}
/*
----------------------------------------------------------
kitamasa法
出典:
m項間漸化式の高速なアルゴリズム - 競技プログラミングをするんだよ
http://nitcoder000.hatenablog.com/entry/kitamasa
一部改変
*/
#define reE(i,a,b) for(auto i=(a);(i)<=(b);i++)
#define rE(i,b) reE(i,0,b)
#define reT(i,a,b) for(auto i=(a);(i)<(b);i++)
#define rT(i,b) reT(i,0,b)
#define rep(i,a,b) reE(i,a,b);
#define rev(i,a,b) for(auto i=(b)-1;(i)>=(a);(i)--)
#define itr(i,b) for(auto i=(b).begin();(i)!=(b).end();++(i))
#define LL int64_t
#define all(b) (b).begin(),(b).end()
/*
使い方
M項間漸化式のN項目を計算する。Nの最小は1。
a[N]=sum(c[i]*a[N-M+i-1])の形。
Mrが本体、Xは使わなくてもできる問題はある。(MODを使わないint,LLなど)
Xは半環を満たすものならなんでもよい、+,*のオーバーロードにMODをねじこむ。
コンストラクタに初項、係数、M,*の単位元、+の単位元の順で引数を与える。
あとはcalcにNを与えるだけ。
*/
#define MAX_LOGN 32
template <class T>
struct Mr{
vector<T> first;
vector<T> C;
vector<vector<T>> bin;
T zero,one;
int M;
//n(1,,,2M)をn(1,,,M)に修正、O(M^2)
void form(vector<T> &n){
rev(i, M + 1, 2 * M + 1){
reE(j, 1, M)n[i - j] = (n[i - j] + (C[M - j] * n[i]));
n[i] = zero;
}
}
//lとrを足し合わせる、O(M^2)
void add(vector<T> &l, vector<T> &r, vector<T> &ans){
reE(i, 1, 2 * M)ans[i] = zero;
reE(i, 1, M)reE(j, 1, M)ans[i + j] = (ans[i + j] + (l[i] * r[j]));
form(ans);
}
//初期化、O(M*MAX_LOGN)
Mr(const vector<T>& f,const vector<T>& c,int m,T e1,T e0){
M = m;
first.reserve(M + 1);C.reserve(M);
zero = e0, one = e1;
first.push_back(zero);
rT(i, M){ first.push_back(f[i]); C.push_back(c[i]); }
bin.resize(MAX_LOGN);
rT(i, MAX_LOGN)bin[i].resize(2*M+1);
rE(i, 2*M)bin[0][i] = zero;
bin[0][1] = one;
reT(i,1, MAX_LOGN){
add(bin[i - 1], bin[i - 1], bin[i]);
}
}
//N項目の計算、戻り値がTの形であることに注意、O(M^2*logN)
T calc(LL n){
n--;
vector<T> tmp,result = bin[0];
for (int b = 0; n; b++,n>>=1)
if (1 & n){ tmp = result; add(tmp, bin[b], result); }
T ans = zero;
reE(i, 1, M)ans = ans + (result[i] * first[i]);
return ans;
}
};
//テンプレート、デフォルトコンストラクタのオーバーロードを忘れない
struct X{
LL val;
X(LL v){ val = v; }
X(){ val = 0; }
LL operator=(const X &another){ return val = another.val; }
LL operator*(const X &another) const {
auto tmp = (val*another.val);
return (tmp + (abs(tmp / MOD) + 1) * MOD) % MOD; // 負のmod対策
}
LL operator+(const X &another)const{ return (val+another.val)%MOD; }
};
int64_t kitamasa(const vector<int64_t>& coefficients, const vector<int64_t>& firsts, int64_t n) {
int m = firsts.size();
vector<X> A, B;
for (auto a : firsts) A.push_back(X(a));
for (auto c : coefficients) B.push_back(X(c));
Mr<X> mr(A, B, m, X(1), X(0));
return mr.calc(n).val;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment