loj 6485 LJJ学二项式定理 —— 单位根反演

题目：https://loj.ac/problem/6485

先把 ( a_{i mod 4} ) 处理掉，其实就是

( sumlimits_{i=0}^{3} a_{i} sumlimits_{j=0}^{n} C_{n}^{j} * s^{j} * [4|(j-i)] )

然后把 ( [4|(j-i)] ) 单位根反演，得到

( sumlimits_{i=0}^{3} a_{i} sumlimits_{j=0}^{n} C_{n}^{j} * s^{j} * frac{1}{4} sumlimits_{k=0}^{3} w_{4}^{k(j-i)} )

( frac{1}{4} sumlimits_{i=0}^{3} a_{i} sumlimits_{k=0}^{3} sumlimits_{j=0}^{n} C_{n}^{j} * s^{j} * w_{4}^{k(j-i)} )

这里 ( w_{4}^{0} = g^{frac{mod-1}{4}} )，( g ) 是 998244353 的原根；

然后为了有二项式定理的形式，构造 ( F(i,w_{4}^{-k}) = w_{4}^{-k(i-n)} ( w_{4}^{-k} + s)^{n} )

于是 ( ans = frac{1}{4} sumlimits_{i=0}^{3} a_{i} sumlimits_{k=0}^{3} F(i,w_{4}^{-k}) )

代码如下：

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
int const mod=998244353,g=3;
ll rd()
{
  ll ret=0,f=1; char ch=getchar();
  while(ch<'0'||ch>'9'){if(ch=='-')f=0; ch=getchar();}
  while(ch>='0'&&ch<='9')ret=ret*10+ch-'0',ch=getchar();
  return f?ret:-ret;
}
ll pw(ll a,ll b)
{
  a=(a%mod+mod)%mod; b=(b%(mod-1)+(mod-1))%(mod-1);
  ll ret=1;
  for(;b;b>>=1,a=a*a%mod)if(b&1)ret=ret*a%mod; return ret;
}
int upt(int x){while(x>=mod)x-=mod; while(x<0)x+=mod; return x;}
int s,a[5]; ll n;
int F(int i,int x)
{
  int ret=upt(x+s); ret=pw(ret,n);
  return (ll)ret*pw(x,i-n)%mod;
}
int main()
{
  int T=rd();
  while(T--)
    {
      n=rd(); s=rd(); for(int i=0;i<4;i++)a[i]=rd();
      int w=pw(g,(mod-1)/4),inv=pw(w,mod-2);
      int ans=0;
      for(int i=0;i<=3;i++)
    for(int k=0,t=1;k<=3;k++,t=(ll)t*inv%mod)
      ans=(ans+(ll)F(i,t)*a[i])%mod;
      ans=(ll)ans*pw(4,mod-2)%mod;
      printf("%d
",upt(ans));
    }
  return 0;
}