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

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

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

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

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

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

\( \frac{1}{4} \sum\limits_{i=0}^{3} a_{i} \sum\limits_{k=0}^{3} \sum\limits_{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} \sum\limits_{i=0}^{3} a_{i} \sum\limits_{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\n",upt(ans));
    }
  return 0;
}