【模板】多项式开根

V.【模板】多项式开根

同之前无数题一样，我们设已知$b^2\equiv A\pmod{x^m}$，并且我们想求出一个$B$使得$B^2\equiv A\pmod{x^{2m}}$。

首先，显然有

$$B-b\equiv0\pmod{x^m}$$

老套路，平方一下，得到

$$B^2-2Bb+b^2\equiv0\pmod{x^{2m}}$$

代入$B^2\equiv A\pmod{x^{2m}}$，有

$$A-2Bb+b^2\equiv0\pmod{x^{2m}}$$

于是移项后得到

$$2Bb\equiv A+b^2\pmod{x^{2m}}$$

除过去，就有

$$B\equiv\dfrac{A+b^2}{2b}\pmod{x^{2m}}$$

然后求解即可。

时间复杂度$O(n\log n)$，但是常数较大，需要O2。

代码：

#include<bits/stdc++.h>
using namespace std;
const int mod=998244353;
const int G=3;
const int N=1<<20;
int n,f[N],g[N],all;
int rev[N];
int ksm(int x,int y){
	int z=1;
	for(;y;y>>=1,x=1ll*x*x%mod)if(y&1)z=1ll*z*x%mod;
	return z;
}
void NTT(int *a,int tp,int LG){
	int lim=(1<<LG),invlim=ksm(lim,mod-2);
	for(int i=0;i<lim;i++)rev[i]=(rev[i>>1]>>1)|((i&1)<<(LG-1));
	for(int i=0;i<lim;i++)if(rev[i]>i)swap(a[i],a[rev[i]]);
	for(int md=1;md<lim;md<<=1){
		int rt=ksm(G,(mod-1)/(md<<1));
		if(tp==-1)rt=ksm(rt,mod-2);
		for(int pos=0,stp=md<<1;pos<lim;pos+=stp){
			int w=1;
			for(int i=0;i<md;i++,w=1ll*w*rt%mod){
				int x=a[pos+i],y=1ll*a[pos+md+i]*w%mod;
				a[pos+i]=(x+y)%mod;
				a[pos+md+i]=(x-y+mod)%mod;
			}
		}
	}
	if(tp==-1)for(int i=0;i<lim;i++)a[i]=1ll*a[i]*invlim%mod;
}
int A[N],B[N],C[N],D[N],E[N];
void mul(int *a,int *b,int *c,int LG){
	int lim=(1<<LG);
	for(int i=0;i<lim;i++)A[i]=B[i]=0;
	for(int i=0;i<(lim>>1);i++)A[i]=a[i],B[i]=b[i];
	NTT(A,1,LG),NTT(B,1,LG);
	for(int i=0;i<lim;i++)A[i]=1ll*A[i]*B[i]%mod;
	NTT(A,-1,LG);
	for(int i=0;i<lim;i++)c[i]=A[i];
}
void inv(int *a,int *b,int LG){
	b[0]=ksm(a[0],mod-2);
	for(int k=1;k<=LG+1;k++){
		mul(a,b,C,k);
		for(int i=0;i<(1<<(k-1));i++)C[i]=(mod-C[i])%mod;
		(C[0]+=2)%=mod;
		mul(C,b,b,k);
	}
}
void sqrt(int *a,int *b,int LG){
	b[0]=a[0];
	for(int k=1;k<=LG+1;k++){
		for(int i=0;i<(1<<(k-1));i++)D[i]=(b[i]<<1)%mod;
		inv(D,E,k-1);
		mul(b,b,b,k);
		for(int i=0;i<(1<<(k-1));i++)(b[i]+=a[i])%=mod;
		mul(b,E,b,k);
	}
}
int main(){
	scanf("%d",&n);
	for(int i=0;i<n;i++)scanf("%d",&f[i]);
	while((1<<all)<n)all++;
	sqrt(f,g,all);
	for(int i=0;i<n;i++)printf("%d ",g[i]);puts("");
	return 0;
}