Kodamanと愉快な仲間たち : T - Powerful Sum

これはまぁ先にTLでヒントを見てしまったので…。
https://www.hackerrank.com/contests/kodamanwithothers/challenges/powerful-sum

問題

2つの実数x,yがあり、和Aと積Bが与えられる。
整数Nが与えられるので、を求めよ。

解法

f(k) = x^k+y^k、g(k)=sum(k(1)+k(2)+....)とする

• f(1) = g(1) = A
• f(2) = A^2 - 2B
• f(k+1) = f(k) * (x+y) - x*y*f(k-1) = A*f(k) - B*f(k-1)
• g(k+1) = g(k) + f(k+1)

なので、あとは漸化式を行列で表現すると(f(k-1),f(k),g(k))→(f(k),f(k+1),g(k+1)の遷移を表現できるので、行列累乗でg(N)を求めればよい。

ll N,A,B;
ll mo=1000000007;

const int MAT=3;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo;
	return r;
}

Mat powmat(ll p,Mat a,int n=MAT) {
	int i,x,y; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(i,n) r.v[i][i]=1;
	while(p) {
		if(p%2) r=mulmat(r,a,n);
		a=mulmat(a,a,n);
		p>>=1;
	}
	return r;
}

ll modpow(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>A>>B;
	A=(A+mo)%mo;
	B=(B+mo)%mo;
	if(N==1) return _P("%d\n",A);
	if(N==2) return _P("%d\n",(A*A+A-2*B+2*mo)%mo);
	
	Mat V;
	V.v[0][0]=V.v[2][1]=1;
	V.v[0][1]=V.v[1][1]=A;
	V.v[0][2]=V.v[1][2]=mo-B;
	V=powmat(N-2,V);
	
	cout<<((A*A+A-2*B+2*mo)%mo*V.v[0][0]+(A*A-2*B+2*mo)%mo*V.v[0][1]+V.v[0][2]*A)%mo<<endl;
	
}

まとめ

以前CFでもうちょいややこしい式変形の問題があったので、それよりはだいぶ易しい。