你是不是瞧不起力扣

leetcode 「10·24」程序员节编程竞赛 计算子集

给你三个整数 $n, k, m$ 。定义 $S=\{i\mid 1\le i\le nm+k,i\in \mathbb Z\}$

请返回一个下标从 $0$ 开始、长度为 $m$ 的数组 answer，其中 answer[i] 表示符合下列条件集合 $T$ 的个数。

集合 $T$ 是集合 $S$ 的子集。
集合 $T$ 中所有元素的和对 $m$ 取余的值恰好为 $i$ 。
由于答案可能很大，你只需要求出它模 998244353 的结果。

0 <= n <= 10^13
0 <= k < min(m, 500)
1 <= m <= 10^5
n * m + k > 0

$k$ 是没有用的求出 $nm$ 的答案再把多出来的那 $k$ 个物品每个暴力 $\mathcal{O}(m)$ 合并进去就行。那么现在就是求：

$$\left(\prod_{i=0}^{m-1}(1+x^i)\right)^n\pmod {x^m-1}$$

循环卷积要做 dft，这里思路 和 UOJ 310 黎明前的巧克力 相同，UOJ 310 是直接把 fwt 写下来，这里来试着写一些 dft，$f_k$ 即为代入 $\omega_m^k$ 的值，这里令 $d=\gcd(k,m)$：

$$\begin{aligned} f_k&=\left(\prod_{i=0}^{m-1}(1+\omega_m^{ik})\right)^n\\ &=\left(\prod_{i=0}^{m/d-1}(1+\omega_{m/d}^{i})\right)^{dn} \end{aligned}$$

这一步是因为 $i$ 模 $m/d$ 相同的 $\omega_m^{ik}$ 值是一样的（单位根的定义），这样化简之后发现指数仅和 $i$ 有关了，括号里面的形式非常好看，考虑比较经典的等式 $\prod\limits_{i=0}^{m-1}(x-\omega_m^i)=x^m-1$（左右两侧均为 $m$ 次多项式且 有 $m$ 处点值相同且 $[x^m]$ 均为 $1$ 故两式相等），将其代入 $x=-1$ 就有：

$$\begin{aligned} &\prod_{i=0}^{m/d-1}(1+\omega_{m/d}^{i})\\ =&(-1)^{m/d}\prod_{i=0}^{m/d-1}(-1-\omega_{m/d}^i)\\ =&(-1)^{m/d}((-1)^{m/d}-1)\\ =&[m/d\bmod 2=1]2 \end{aligned}$$

于是 $f_k=[m/d\bmod 2=1]2^{dn}$，这里发现 $f_k$ 的值仅和 $d$ 有关，那么在 idft 的时候就可以利用这个转成枚举 $d$：

$$\begin{aligned} g_k=&\frac{1}{m}\sum_{i=0}^{m-1}f_k\omega_m^{-ik}\\ =&\frac{1}{m}\sum_{d\mid m}f_d\sum_{i=0}^{m/d-1}[\gcd(i,m/d)=1]\omega_{m/d}^{-ik}\\ =&\frac{1}{m}\sum_{d\mid m}f_d\sum_{i=0}^{m/d-1}\sum_{e\mid i,e\mid m/d}\mu(e)\omega_{m/d}^{-ik}\\ =&\frac{1}{m}\sum_{de\mid m}f_d\mu(e)\sum_{0\leq i<m/de}\omega_{m/de}^{-ik}\\ =&\frac{1}{m}\sum_{de\mid m}f_d\mu(e)\frac{m}{de}[\frac{m}{de}\mid k]\\ \end{aligned}$$

我朴素实现了个 $\mathcal{O}(m(\log m+d(m)))$。

代码

class Solution {
public:
	typedef long long ll;
	typedef vector<int>vi;
	#define pb emplace_back
	static const int mod=998244353;
    static const int N=100010;
	inline void cadd(int &x,int y){x=(x+y>=mod)?(x+y-mod):(x+y);}
	inline void cdel(int &x,int y){x=(x-y<0)?(x-y+mod):(x-y);}
	inline int add(int x,int y){return (x+y>=mod)?(x+y-mod):(x+y);}
	inline int del(int x,int y){return (x-y<0)?(x-y+mod):(x-y);}
	int qpow(int x,ll y){
		int s=1;
		while(y){
			if(y&1)s=1ll*s*x%mod;
			x=1ll*x*x%mod;
			y>>=1;
		}
		return s;
	}
	ll n,m,k;
	int vis[N];
	vi pr;
	int f[N],g[N];
	vi vec[N];
	void init(){
		for(int i=2;i<=m;i++){
			if(!vis[i])
				pr.pb(i);
			for(auto j:pr){
				if(i*j>m)break;
				vis[i*j]=1;
				if(i%j==0)break;
			}
		}
	}
    ll gcd(ll x,ll y){return !y?x:gcd(y,x%y);}
    vector<int> subsetCounting(long long q, int w, int e) {
    	n=q;k=w;m=e;
		if(n){
			init();
			for(int i=1;i<=m;i++){
				int d=__gcd(1ll*i,m);
				if(m/d%2==1){
					f[i]=qpow(2,1ll*d*n);
				}
			}
			for(auto i:pr){
				for(int j=m/i;j;j--)
					cdel(f[i*j],f[j]);
			}
			for(int i=1;i<=m;i++)if(m%i==0)vec[m].pb(i);
			int iv=qpow(m,mod-2);
			for(int i=1;i<=m;i++){
				for(auto de:vec[m]){
					if(i%(m/de)==0){
						cadd(g[i],1ll*f[de]*(m/de)%mod);
					}
				}
				g[i]=1ll*g[i]*iv%mod;
			}
			g[0]=g[m];
			g[m]=0;
		}
		else g[0]=1;
		for(int i=0;i<=m;i++)f[i]=0;
		for(int i=1;i<=k;i++){
			for(int j=0;j<m;j++)
				cadd(f[(j+i)%m],g[j]);
			for(int j=0;j<m;j++)
				g[j]=add(f[j],g[j]),f[j]=0;
		}
		vi vec;
		for(int i=0;i<m;i++)vec.pb(g[i]);
		return vec;
    }
};