LOJ#6713. 「EC Final 2019」狄利克雷 k 次根 加強版

題目描述

定義兩個函數 \(f, g: \{1, 2, \dots, n\} \rightarrow \mathbb Z\) 的狄利克雷卷積 \(f * g\) 為：

\[(f * g)(n) = \sum_{d | n} f(d)g(\frac nd) \]

我們定義 \(g = f^k\) 即 \(k\) 次冪為：

\[f^{k}=\underbrace {f * \dots * f} _{k~{\textrm {個}}} \]

在本題中，我們想要解決這個問題的逆問題：給你 \(g\) 和 \(k\)，你需要找到一個函數 \(f\) 使得 \(g = f^k\)。

另外，保證 \(g(1)=1\)，你需要保證 \(f(1)=1\)。所有的運算在 \(\mathbb F_p\) 上進行，其中 \(p = 998244353\)，這意味着狄利克雷卷積為 \((f*g)(n) = \left(\sum_{d|n} f(d)g(\frac nd)\right) \bmod p\)。

\(n\le 10^6\)

Sol

定義一個數論函數 \(f(n)\) 的狄利克雷生成函數是 \(F(x)=\sum_{n=1}^\infty f(n)n^x\)。

注意這個東西不是關於 \(x\) 的多項式（一開始我當成了多項式），然而它可以做多項式乘法，由於這個 \(a^x\cdot b^x \rightarrow (ab)^x\)，所以對於生成函數系數來說，相當於完成了兩個數論函數的 狄利克雷卷積。

既然有類似多項式的性質，那么我們考慮 ln-exp 的方法求解 k 次根。

根據多項式 ln 的推導，有 \(B=\ln A\rightarrow B=\int {A'\over A}\)

除法可以用 狄利克雷除法 實現，那么我們只需要支持 導數和積分 運算即可。

考慮這個指數函數的求導法則 \(\displaystyle {\text{d} a^x\over \text{d}x}=a^x\ln a\)。

然而在正整數意義下沒有 \(\ln\) 這種運算，然后做法是找一個運算性質跟 \(\ln\) 一樣的東西來替代它，比如質因子個數函數 \(\text{cnt}_x\)（相同的算作多個），因為它滿足 \(f(x)+f(y)=f(xy)\) 這樣的性質。

於是可以實現一個對一個狄利克雷級數求對數函數的代碼如下：

void ln(int n, int *a)
{
	b[1] = 0;
	for(int i = 2; i <= n; ++i)
		b[i] = (ll)a[i] * cnt[i] % P;
	for(int i = 2; i <= n; ++i)
	{
		for(int j = 2, k = i * j; k <= n; ++j, k += i)
			b[k] = (b[k] - 1LL * b[i] * a[j]) % P;
		b[i] = (1LL * b[i] * inv[cnt[i]] % P + P) % P;
	}
	for(int i = 1; i <= n; ++i)
		a[i] = b[i];
}

由於 \(B=A^{1/k}=\exp {{\ln A} \over k}\)，那么只要先求 \(\ln\) 再 \(\exp\) 即可。
求指數函數也是類似的，整個代碼如下，和暴力 \(O(n^2)\) 的普通多項式對數/指數函數非常相似。

#include<stdio.h>

typedef long long ll;
const int N = 1000005, P = 998244353;

int fexp(int a, int b)
{
	ll x = 1, o = a;
	for(; b; b >>= 1, o = o * o % P)
		if(b & 1) x = x * o % P;
	return x;
}

int n, K, a[N], b[N], inv[128], cnt[N];
bool np[N];

void init()
{
	cnt[1] = 1;
	for(int i = 2; i <= n; ++i)
		if(!np[i])
		{
			cnt[i] = 1;
			for(int j = 2, k = i * j; k <= n; j++, k += i)
			{
				np[k] = true;
				if(!cnt[k] && cnt[j]) cnt[k] = cnt[j] + 1;
			}
		}
	inv[1] = 1;
	for(int i = 2; i < 128; ++i)
		inv[i] = (ll)inv[P % i] * (P - P / i) % P;
}

void ln(int n, int *a)
{
	b[1] = 0;
	for(int i = 2; i <= n; ++i)
		b[i] = (ll)a[i] * cnt[i] % P;
	for(int i = 2; i <= n; ++i)
	{
		for(int j = 2, k = i * j; k <= n; ++j, k += i)
			b[k] = (b[k] - 1LL * b[i] * a[j]) % P;
		b[i] = (1LL * b[i] * inv[cnt[i]] % P + P) % P;
	}
	for(int i = 1; i <= n; ++i)
		a[i] = b[i];
}

void exp(int n, int *a)
{
	for(int i = 2; i <= n; ++i)
		b[i] = (ll)a[i] * cnt[i] % P, a[i] = 0;
	a[1] = 1;
	for(int i = 1; i <= n; ++i)
	{
		a[i] = 1ll * a[i] * inv[cnt[i]] % P;
		for(int j = 2, k = i * j; k <= n; ++j, k += i)
			a[k] = (a[k] + 1LL * a[i] * b[j]) % P;
	}
}

int main()
{
	scanf("%d %d", &n, &K);
	for(int i = 1; i <= n; ++i)
		scanf("%d", a + i);
	init();
	K = fexp(K, P - 2);
	ln(n, a);
	for(int i = 1; i <= n; ++i) a[i] = 1ll * a[i] * K % P;
	exp(n, a);
	for(int i = 1; i <= n; ++i) printf("%d%c", a[i], " \n"[i == n]);
	return 0;
}