kitamasa法でm項間漸化式の第n項までの和を求めるやつ

nth_sum_kitamasa.cpp

#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;
}