[LOJ 6485] LJJ binomial theorem (unit root inversion) (template)

LJJ Learn binomial theorem

Topic link ：LOJ 6485

The main idea of the topic

Find an equation ：
$\sum _{i=0}^n((\genfrac{}{}{0ex}{}{n}{i})s^i{a}_{i\mathrm{m}\mathrm{o}\mathrm{d}4})$
among $n\le 10^{18}$

Ideas

This problem can be regarded as the template problem of unit root inversion .

Unit root inversion

First, the unit root inversion is used to deal with the formula with modulus .
（ There is a module in the coefficient that can be removed with this ）

Then the formula is $[n|a]=\frac{1}{n}\sum _{k=0}^{n-1}{\omega }_n^{ak}$

Then prove it ：
if $a\ne 0(\mathrm{m}\mathrm{o}\mathrm{d}n)$
Then we can sum it with an equal ratio sequence ：
$\frac{1}{n}\frac{{\omega }_n^{an}-1}{{\omega }_n^a-1}$
then ${\omega }_n^a\ne 1,{\omega }_n^a-1\ne 0$ Denominator is ok , Molecular cause ${\omega }_n^n={\omega }_n^0=1$,${\omega }_n^{an}=({\omega }_n^n{)}^a=1,{\omega }_n^{an}-1=0$

if $a=0(\mathrm{m}\mathrm{o}\mathrm{d}n)$
That special treatment （ Because at this time, there is $0$）
$\frac{1}{n}\sum _{k=0}^{n-1}{\omega }_n^{ak\mathrm{m}\mathrm{o}\mathrm{d}n}$
$\frac{1}{n}\sum _{k=0}^{n-1}{\omega }_n^0=\frac{1}{n}n=1$

Then a classic formula ：
$[a=b(\mathrm{m}\mathrm{o}\mathrm{d}n)]=[a-b=0(\mathrm{m}\mathrm{o}\mathrm{d}n)]=[n|(a-b)]$
$=\frac{1}{n}\sum _{k=0}^{n-1}{\omega }_n^{(a-b)k}=\frac{1}{n}\sum _{k=0}^{n-1}{\omega }_n^{ak}{\omega }_n^{-bk}$

This question

$\sum _{i=0}^n((\genfrac{}{}{0ex}{}{n}{i})s^i{a}_{i\mathrm{m}\mathrm{o}\mathrm{d}4})$
$\sum _{i=0}^n((\genfrac{}{}{0ex}{}{n}{i})s^i(\sum _{j=0}^3{a}_j[i\mathrm{m}\mathrm{o}\mathrm{d}4=j]))$
$\sum _{i=0}^n((\genfrac{}{}{0ex}{}{n}{i})s^i(\sum _{j=0}^3{a}_j[4|i-j]))$
$\sum _{i=0}^n((\genfrac{}{}{0ex}{}{n}{i})s^i(\sum _{j=0}^3{a}_j(\frac{1}{4}\sum _{k=0}^3{\omega }_4^{ik}{\omega }_4^{-jk})))$
$\frac{1}{4}\sum _{k=0}^3(\sum _{j=0}^3{a}_j{\omega }_4^{-jk})(\sum _{i=0}^n(\genfrac{}{}{0ex}{}{n}{i})s^i{\omega }_4^{ik})$

（ binomial theorem ）
$\sum _{i=0}^n\left(\genfrac{}{}{0ex}{}{n}{i}\right)s^i{\omega }_4^{ik}$
$\sum _{i=0}^n\left(\genfrac{}{}{0ex}{}{n}{i}\right)(s{\omega }_4^k{)}^i{1}^{n-i}$
$(s{\omega }_4^k+1{)}^n$

（ Back to the original ）
$\frac{1}{4}\sum _{k=0}^3(\sum _{j=0}^3{a}_j{\omega }_4^{-jk})(s{\omega }_4^k+1{)}^n$

then ${\omega }_4^1=g^{(mod-1)/4}$.

Code

#include<cstdio>
#define ll long long
#define mo 998244353

using namespace std;

ll n, s, a[4], G = 3, Gv, v4, ans, w[4];

ll ksm(ll x, ll y) {
    ll re = 1; while (y) {
    if (y & 1) re = re * x % mo; x = x * x % mo; y >>= 1;} return re;}

int main() {
    
	int T; scanf("%d", &T); Gv = ksm(G, mo - 2); v4 = ksm(4, mo - 2);
	w[0] = 1; w[1] = ksm(G, (mo - 1) / 4); for (int i = 2; i < 4; i++) w[i] = w[i - 1] * w[1] % mo;
	while (T--) {
    
		scanf("%lld %lld", &n, &s); for (int i = 0; i < 4; i++) scanf("%lld", &a[i]);
		ans = 0;
		for (int k = 0; k <= 3; k++) {
    
			ll now = 0;
			for (int j = 0; j <= 3; j++)
				(now += a[j] * w[(4 - j) * k % 4] % mo) %= mo;
			(ans += now * ksm((s * w[k] % mo + 1) % mo, n) % mo) %= mo;
		}
		printf("%lld\n", ans * v4 % mo);
	}
	
	return 0;
}