Official

G - Necklace Editorial by Nyaan


次の 2 通りの解法を解説します。

  1. Polya の数え上げ定理 を用いた初等的な \(\mathrm{O}(U \log^2 U)\) 解法
  2. ディリクレ母関数 を用いた高度な \(\mathrm{O}(U \log U)\) 解法

2 番目のディリクレ母関数を用いた解法では、魔法としか言いようのないテクニックが登場します。本当にビックリするアルゴリズムなので是非解説を読んでみてください。

1. Polya の数え上げ定理を用いた解法

Polya の数え上げ定理については順列の場合を ABC284-Ex の解説で詳説したのでここでは説明を省略します。

まず、はじめに列の場合を解いてみましょう。\(A(n)\) を次のように定義します:長さ \(U\) の数列であって、\(i\) 番目の要素 \(A(n)[i]\) が「総積が \(i\) であるような \(n\) 個の宝石の列としてありうるものの個数」である。

制約より、列の長さが \(\lfloor \log_2 U \rfloor\) より長いものは考えなくてよいです。\(V = \lfloor \log_2 U \rfloor\) と置きます。
この時、\(A(0), A(1), \dots, A(V)\) は以下の疑似コードで表される DP で \(\mathrm{O}(U \log^2 U)\) で容易に求まります。

using mint = atcoder::modint998244353;
int freq[500003]; // b_i の頻度列, 計算済
mint A[21][500003];
for(int i = 1; i <= U; i++) A[0][i] = (i == 1 ? 1 : 0);
for(int i = 1; i <= V; i++)
  for(int s = 1; s <= U; s++)
    for(int t = 1; s * t <= U; t++)
      A[i + 1][s * t] += A[i][s] * freq[t];

計算で得られた \(A(n)\) を用いて円環の場合を考えてみましょう。
長さ \(L\) で総積が \(x\) のネックレスの個数をポリアの数え上げ定理を用いて計算しましょう。不動点の個数の平均を取ってみると、

\[ \begin{aligned} &\frac{1}{L} \sum_{i=0}^{L-1} A\left( \frac{L}{\gcd(i, L)}\right) [\sqrt[\gcd(i,L)]{x}] \\ &= \frac{1}{L} \sum_{d \vert L} \varphi(d) A \left( \frac{L}{d} \right) [\sqrt[d]{x}] \end{aligned} \]

となります。ただし \(\sqrt[d]{x}\) が整数値を取らない時は \(\sqrt[d]{x}\) を添え字に含む項は \(0\) とみなします。

列の時と同様、\(L\)\(V\) 以下の範囲を考えればよいとわかります。よって上式を \(L=1,2,\dots,V\) の範囲で足し合わせれば総積が \(x\) のネックレスの個数を計算することが出来て、これはトーシェント関数などを適切に前計算していれば \(\mathrm{O}(V \log V)\) で計算することが出来ます。

以上よりこの問題を \(\mathrm{O}(U \log^2 U)\) で解くことができて、これは十分高速です。

ディリクレ母関数を用いた解法

ディリクレ母関数とは、数列 \(a_1, a_2, \dots\) に対して次の式で表される母関数です。

\[A(s) = \frac{a_1}{1^s} + \frac{a_2}{2^s} + \dots = \sum_{n=1}^\infty \frac{a_n}{n^s}\]

ディリクレ母関数は今回のような総積が重要な問題にうってつけの母関数です。

「美しさ \(2\) の宝石」と「美しさ \(3\) の宝石」を組み合わせると「美しさの総積が \(6\) の宝石の列」になるという変化は

\[\frac{1}{2^s} \times \frac{1}{3^s} = \frac{1}{6^s}\]

のようにディリクレ母関数の積に対応します。

さて、ディリクレ母関数を用いて今回の問題を記述してみましょう。 宝石の美しさの頻度列のディリクレ母関数を \(A(s)\) と置きます。

長さが \(L\) のネックレスの母関数 \(C_L(s)\) を Polya の数え上げ定理を用いて計算してみると次式になります。

\[C_L(s) = \frac{1}{L} \sum_{D \vert L} \varphi(d) A(ds)^{L/d}\]

よって、ネックレスの母関数 \(C(s)\)

\[ \begin{aligned} C(s) &= \sum_L C_L(s) \\ &= \sum_L \frac{1}{L} \sum_{D \vert L} \varphi(d) A(ds)^{L/d} \\ &= \sum_{D \vert L}\frac{1}{L} \varphi(d) A(ds)^{L/d} \\ &= \sum_{D} \sum_{t \geq 1} \frac{1}{dt} \varphi(d) A(ds)^{t} \\ &= \sum_{D} \frac{\varphi(d)}{d}\sum_{t \geq 1} \frac{A(ds)^{t} }{t} \\ &= \sum_{D} \frac{\varphi(d)}{d} \log \left( \frac{1}{1 - A(ds)}\right) \end{aligned} \]

となり、この式を計算できればよいです。ここで \(\log \left( \frac{1}{1 - A(2s)}\right)\)\(\log \left( \frac{1}{1 - A(3s)}\right)\) のような項は \(\log \left( \frac{1}{1 - A(s)}\right)\)\(\frac{c}{n^s}\) のような項を \(\frac{c}{(n^2)^s}\)\(\frac{c}{(n^3)^s}\) にしたものだとわかるので \(\log \left( \frac{1}{1 - A(s)}\right)\) がわかっていれば全て計算できます。また、 \(1 \leq D \leq \lfloor \log_2 U \rfloor\) の範囲だけ調べればよいことも明らかです。
よって \(\log \left( \frac{1}{1 - A(s)}\right)\) が計算できれば今回の問題が解けることがわかりました。

\(U\) 項目で打ち切ったディリクレ母関数の畳み込みや逆元が \(\mathrm{O}(U \log U)\) で計算できることは容易に確認できますから(わからない人は考えてみましょう)、\(\log\) の計算が目下の課題となります。
一法としてはテイラー展開すると \(\lfloor \log_2 U \rfloor + 1\) 項以降が全て \(0\) になることから、前 \(\lfloor \log_2 U \rfloor\) 項を畳み込みで計算するというのがあり、これで \(\mathrm{O}(U \log^2 U)\) 解法が従います。
より高速に計算するために微分を使ってみましょう。 既知の関数 \(G(s)\) から未知の関数 \(F(s) = \log G(s)\) を求めることを考えます。

\[ \begin{aligned} &\log G = F \\ &\to e^F = G \\ &\to F' e^F = F' G = G'\\ &\to F' = \frac{G'}{G} \end{aligned} \]

となり微積分とディリクレ母関数の畳み込み・逆元を組合せれば \(G\) から \(F\) を計算できそうな気分になります。
しかしここで問題が発生します。\(a_1, a_2, \) のディリクレ母関数 \(A(s)\) の微分は

\[A'(s) = -\sum_{n =1}^\infty \frac{a_n \ln n}{n^s}\]

となり自然対数 \(\ln n\) が出てきてしまいました。とうぜん \(\ln n\) は mod 998244353 上で表せないため、微分の計算における大きな障害となります。

困ってしまったかのようですがここで魔法が登場します。発見者は apiad のブログによると _rqy と Elegia とのことです。
さて、その魔法の名は \(n\) の重複も含めた素因数の個数の関数 \(\Omega(n)\) です。砕けた言い方で説明すると \(1\) に素数を上手く掛け続けて \(n\) になるまでにかかる回数を意味する関数と言えます。 一見よくある関数にしかみえませんが、実は \(\Omega(n)\)\(-\ln n\) と置き換えることができるのです!

形式的に説明します。特殊な微分演算子 \(\mathfrak{D}\) を次のように定義します。

\[\mathfrak{D} A(s) = \sum_{n =1}^\infty \frac{a_n \Omega(n)}{n^s}\]

すると、驚くべきことにディリクレ母関数 \(F, G\) および通常の関数 \(h\) について

\[\mathfrak{D} (FG) = F \mathfrak{D} G + G \mathfrak{D} F\]

\[\mathfrak{D}(h \circ G) = (h' \circ G) \mathfrak{D} G\]

が成り立つことが証明できます。つまりこの魔法の微分は通常の微積分と同様に扱って良いのです!

よって、\(F = \log(G)\) の計算は

\[\mathfrak{D} F = \frac{ \mathfrak{D} G}{G}\]

を満たす \(F\) を計算すればよいことになり、これは \(\mathrm{O}(U \log U)\) で計算できます。 このように、ディリクレ母関数の (mod 998244353 上の) 微分は定義できませんが、少し微分の定義を変えれば全てが上手く回るようになる、という魔法のようなテクニックの紹介でした。
以上の内容を適切に実装すればすべてを \(\mathrm{O}(U \log U)\) で計算することが出来て、これは非常に高速です。

実装例 (Python, codon) は次の通りです。言語アップデート回とのことなので codon で実装してみました。かなり雑な一発書きでも 120 ms 程度で動作しています。Python で modint が十分高速に動作するのは非常に気持ちが良いですね。
さて、2025 年現在主流の Python 処理系である PyPy は苦手分野が数知れず存在して (再帰, tuple の sort, データ構造全般, modint, 抽象化…) 、それらがアルゴという競技における Python の存在の歪みを生み出すとともに、Python 1 本で戦うプレイヤーにとっての足かせとなっていることはこの解説を読む層であればよくご存じかと思います。codon はこれら全てを解決してくれる処理系とのことで、功罪半ばする PyPy を押しのけて tier 1 の処理系に成り上がることを強く熱望しています。(CPython との細かい仕様の差異を吸収するテンプレートが早く作られないかな、特に map の挙動周り。あと ACL for codon もほしい)

@tuple
class modint:
    x: u32

    def __new__():
        return modint(u32(0))

    def __new__(y: int) -> modint:
        z = y % 998244353
        if z < 0:
            z += 998244353
        return modint(u32(z))

    def __add__(self, rhs: modint) -> modint:
        res = self.x + rhs.x
        if res >= u32(998244353):
            res -= u32(998244353)
        return modint(res)

    def __sub__(self, rhs: modint) -> modint:
        res = self.x - rhs.x
        if res >= u32(998244353):
            res += u32(998244353)
        return modint(res)

    def __mul__(self, rhs: modint) -> modint:
        res = u64(self.x) * u64(rhs.x) % u64(998244353)
        return modint(u32(res))
    
    def __neg__(self) -> modint:
        return modint(u32(0) if self.x == u32(0) else u32(998244353) - self.x)

    def __str__(self) -> str:
        return str(self.x)

def dirichlet_conv(a, b):
    c = [modint(0)] * (U + 1)
    for i in range(1, U + 1):
        for j in range(1, U // i + 1):
            c[i * j] += a[i] * b[j]
    return c

def dirichlet_inv(a):
    assert a[1] == modint(1)
    b = [modint(0)] * (U + 1)
    b[1] = modint(1)
    for i in range(1, U + 1):
        for j in range(2, U // i + 1):
            b[i * j] -= b[i] * a[j]
    return b

N, U = [int(a) for a in input().split()]

# Omega(n)
lpf = [i for i in range(U + 1)]
for p in range(2, U + 1):
    if lpf[p] != p:
        continue
    for j in range(p * p, U + 1, p):
        lpf[j] = min(lpf[j], p)
Omega = [0] * (U + 1)
for i in range(2, U + 1):
    Omega[i] = Omega[i // lpf[i]] + 1

# totient
totient = [i for i in range(30)]
for i in range(1, 30):
    for j in range(2 * i, 30, i):
        totient[j] -= totient[i]

# modinv
inv = [modint(0)] + [modint(pow(i, 998244353 - 2, 998244353)) for i in range(1, 30)]

A = [modint(0)] * (U + 1)
for b in map(int, input().split()):
    A[b] += modint(1)

# 1 - A(s)
B = [modint(0), modint(1)] + [-x for x in A[2:]]
# 1 / (1 - A(s)) = inv(B)
G = dirichlet_inv(B)
# log(G) = F <-> F' = G' B
for i in range(1, U + 1):
    G[i] *= modint(Omega[i])
F = dirichlet_conv(G, B)
for i in range(2, U + 1):
    F[i] *= inv[Omega[i]]

ans = [modint(0)] * (U + 1)
for i in range(2, U + 1):
    j, d = i, 1
    while j <= U:
        ans[j] += F[i] * modint(totient[d]) * inv[d]
        j, d = j * i, d + 1
print(" ".join(map(str, ans[2:])))

posted:
last update: