G - Kyoen 解説 by en_translator

より高速な解法

公式解説 で説明された知識と \(O(MC^4 + P^{1.5})\) 解法を前提に、より高速な解法を紹介します。

1. 本問題の解法のおさらい

本問題は概ね以下の問題に帰着されます。表記の簡略化のため、\(\sqrt{-1}\) を \(i\) と表記します。
（元の問題からは \(A, B\) の符号を反転させていることに注意してください。）

整数 \(A, B, C\)、長さ \(M+1\) の整数列 \((Q_0, Q_1, \ldots, Q_M), (R_0, R_1, \ldots, R_M)\)、長さ \(M\) の非負整数列 \((E_1, \ldots, E_M)\) が与えられる。 ただし、\(Q_j^2 + R_j^2\) \((j = 1, \ldots, M)\) は相異なる有理素数である。 非負整数列 \((k_0, k_1, \ldots, k_M)\) (\(k_0 < 4\), \(k_j \leq E_j\) (\(1 \leq j \leq M\))) であって、 \[ i^{k _ 0} (Q_0 + i R_0) \prod _ {j=1}^M (Q_j + i R_j)^{k _ j} (Q _ j - i R _ j)^{E _ j - k _ j} \] と表されるガウス整数の実部と虚部を \(C\) で割った余りがそれぞれ \(A\) および \(B\) となるようなものの個数を求めよ。

解説で紹介された解法は概ね以下と等価です。ここで、ガウス整数の演算はすべて \(\bmod\) \(C\) で考えます。

  • 二次元配列 \(\mathrm{dp}[x][y]\) （\(0 \leq x, y < C\)、以下同） を \(0\) で初期化する。
    • \(\mathrm{dp}[1 \bmod C][0]\) を \(1\) で初期化する。
  • \(j = 1, \ldots, M\) に対して以下を行う。
    • 二次元配列 \(\mathrm{cnt}[x][y]\) を、\[ x + yi = (Q_j + i R_j)^{k_j} (Q_j - i R_j)^{E_j - k_j} \quad \ldots(\star) \] となるような \(k_j\) の個数とする。
    • 二次元配列 \(\mathrm{dp}'[x][y]\) のすべての要素を \(0\) で初期化する。
      • 各 \(x_1, y_1, x_2, y_2\) に対して、\((x + yi) = (x_1 + i y_1)(x_2 + i y_2)\) として、\(\mathrm{dp'}[x][y]\) に \(\mathrm{dp}[x_1][y_1] \cdot \mathrm{cnt}[x_2][y_2]\) を加算する。
    • \(\mathrm{dp} \gets \mathrm{dp}'\) と更新する。
  • \(k_0 = 0, \ldots, 3\) について、\(A' + iB' = i^{k_0} (Q_0 + i R_0) (A + iB)\) を求め、答えに \(\mathrm{dp}[A'][B']\) を加算する。

\(\mathrm{dp}[x][y]\) と \(\mathrm{cnt}[x][y]\) の要素数はいずれも \(C^2\) です。
また、\((\star)\) が最初と末尾の \(C^2\) 項を除くと周期 \(C^2\) 以下の周期性を持つことから、 \(\mathrm{cnt}\) は計算量 \(O(C^2)\) で求められます。
したがって、このアルゴリズムの計算量は \(O(M C^4)\) となります。
（元の問題を上記の問題に帰着させる際、\(Q_i, R_i\) を求めるのに \(O(P^{1.5})\) 時間かかることから、全体の計算量が \(O(M C^4 + P^{1.5})\) となっています。）

2. \(O(M C^{2 + 2 (\log 2) / (\log \log C)} + P^{1.5}) = O(M C^{2+\epsilon} + P^{1.5})\) 解法

実は、\(\mathrm{dp}\) および \(\mathrm{cnt}\) の取りうる添字 \((x, y)\) の個数は高々 \(C\cdot d( C )\) 個であることが示せます。 ここで、\(d(n)\) は正整数 \(n\) の正の約数の個数です。

2-1. 証明

  • 正の整数 \(n\) に対して、その素因数の集合を \(\mathcal P(n)\) と書きます。
  • 正の整数 \(n, m\) に対して、\(m\) が \(n\) を割り切る回数を \(\nu_m(n)\) と書きます。
    • また、\(\nu_m(0) = \infty\) と定義します。
  • 正の整数 \(n\) に対して、実部と虚部を \(\bmod\) \(n\) で考えたときのガウス整数の集合を \(\mathbb{Z}/n\mathbb{Z}[i]\) と書きます。
    • そのうち積に関する可逆元からなる集合（乗法群）を \((\mathbb{Z}/n\mathbb{Z}[i])^\times\) と書きます。

2-1-1. 任意の整数 \(x, y\) に対して、\((x + yi) \in (\mathbb{Z}/C\mathbb{Z}[i])^\times \iff \gcd(x^2 + y^2, C) = 1\) を示します。

  • \((\Leftarrow)\) \(\gcd(x^2 + y^2, C) = 1\) ならば、\(\bmod\) \(C\) での逆元 \((x^2 + y^2)^{-1}\) が取れます。 これを実部と虚部に掛けた \((x^2 + y^2)^{-1}x + (x^2 + y^2)^{-1}yi\) が逆元となっていることが確かめられます。
  • \((\Rightarrow)\) 対偶を示します。 \(g \coloneqq \gcd(x^2 + y^2, C) > 1\) より、整数 \(g' \coloneqq (C/g)\) は \(1 \leq g' < C\) を満たします。
    \((x^2 + y^2) g' \equiv 0 \pmod C\) より、\(\mathbb Z/C\mathbb Z[i]\) 上で \((x + yi)(x - yi)(g' + 0i) = (0 + 0i)\) です。
    もし \(xg' \equiv yg' \equiv 0 \pmod C\) なら \((x + yi) (g' + 0i) = (0 + 0i)\) であり、それ以外の場合は \((x + yi) (xg' - yg'i) = (0 + 0i)\) です。
    よって、\(z \in \mathbb Z/C\mathbb Z[i]\) を \(xg' \equiv yg' \equiv 0 \pmod C\) なら \(z = g' + 0i\)、それ以外の場合は \(z = xg' - yg'i\) とおくと、 \(z \neq 0 + 0i\) かつ \((x + yi) z = 0 + 0i\) となります（すなわち、\(z\) は零因子となります）。
    しかし、\((x + yi)^{-1}\) が存在するならば、\(0 + 0i = (x + yi)^{-1} (0 + 0i) = (x + yi)^{-1} (x + yi) z = z\) となり、矛盾します。
    したがって、\((x + yi)^{-1}\) は存在しません。

2-1-2. \(\mathbb{Z}/C\mathbb{Z}[i]\) のうち「逆元をもつもの」の個数 \(f(C)\) が、 素数 \(p\) と正の整数 \(e\) に対して \[ f(p^e) = \begin{cases} 2^{2e-1} & (p = 2) \\ (p^{e-1})^2 (p-1)^2 & (p \equiv 1 \pmod 4) \\ (p^{e-1})^2 (p^2-1) & (p \equiv 3 \pmod 4) \end{cases} \] を満たす乗法的関数であることを示します。
\(f(C) = |(\mathbb Z/C\mathbb Z[i])^\times|\) は「整数組 \((x, y)\) であって、\(0 \leq x, y < C\) かつ \(\gcd(x^2 + y^2, C) = 1\) を満たすものの個数」に等しいです。
中国剰余定理から、この個数は各 \(p \in \mathcal P(C)\) についての「整数組 \((x, y)\) であって、\(0 \leq x, y < p^{\nu_p(C)}\) かつ \(\gcd(x^2 + y^2, C) = 1\) を満たすものの個数」の積に等しいです。
\(\gcd(x^2 + y^2, C) = 1 \iff \gcd(x^2+y^2, p^e) = 1\) から、\(\displaystyle f(C) = \prod_{p\in\mathcal P(C)} f(p^{\nu_p(C)})\)、すなわち \(f\) が乗法的関数であることが示せます。

素数 \(p\) と正整数 \(e\) に対して、 \(f(p^e)\) は「整数組 \((x, y)\) であって、\(0 \leq x, y < p^e\) かつ \(\gcd(x^2+y^2, p^e)=1 \iff x^2 + y^2 \not\equiv 0 \pmod p\) を満たすものの個数」であり、 これは「整数組 \((x, y)\) であって、\(0 \leq x, y < p\) かつ \(x^2 + y^2 \not\equiv 0 \pmod p\) を満たすものの個数」すなわち \(f(p)\) の \((p^{e-1})^2\) 倍に等しいです。

\(p = 2\) のとき、\((x, y) = (0, 1), (1, 0)\) なので、\(f(2) = 2\) です。以下、\(p\) は奇素数とします。

補集合、すなわち「整数組 \((x, y)\) であって、\(0 \leq x, y < p\) かつ \(x^2 + y^2 \equiv 0 \pmod p\) を満たすもの」の個数を求めます。 \(x = 0\) のとき \(y^2 \equiv 0 \pmod p\) すなわち \(y = 0\) です。以下、\(x > 0\) とします。 すると、\(x\) の \(\bmod\) \(p\) での逆元 \(x^{-1} \neq 0\) が取れるので、\(x^2 + y^2 \equiv 0 \iff (x^{-1} y)^2 \equiv -1 \pmod p\) がわかります。 すなわち、それぞれの \(x\) と \(s^2 \equiv -1 \pmod p\) なるそれぞれの \(s\) に対して、\(y = x^{-1}s\) が条件を満たします。

平方剰余に関する有名な結果から、 \(p \equiv 1 \pmod 4\) のとき \(s\) はちょうど \(2\) 個の解を、\(p \equiv 3 \pmod 4\) のとき \(s\) はちょうど \(0\) 個の解を持ちます。 したがって補集合の要素数は \(p \equiv 1 \pmod 4\) のとき \(1 + (p-1) \cdot 2\) 個であり、\(p \equiv 3 \pmod 4\) のとき \(1\) 個です。 よって \[ f(p) = \begin{cases} (p - 1)^2 & (p \equiv 1 \pmod 4) \\ p^2 - 1 & (p \equiv 3 \pmod 4) \end{cases} \] となって、与式が示されます。

2-1-3. 逆元を持つ \((a + bi) \in \mathbb{Z}/C\mathbb{Z}[i]\) に対して、\(\dfrac{a+bi}{a-bi}\) が取りうる値の種類数 \(g(C)\) が、 素数 \(p\) と正の整数 \(e\) に対して \[ g(p^e) = \begin{cases} p^{e-1} (p-1) & (p \not\equiv 3 \pmod 4) \\ p^{e-1} (p+1) & (p \equiv 3 \pmod 4) \end{cases} \] を満たす乗法的関数であることを示します。
写像 \(\psi_C\colon (\mathbb{Z}/C\mathbb{Z}[i])^\times \to (\mathbb{Z}/C\mathbb{Z}[i])^\times\) を \(\psi_C(a+bi) \coloneqq \dfrac{a+bi}{a-bi}\) で定義すると、 これは群の準同型写像であることが示せます。したがって、群の第一同型定理により、 \[ | \operatorname{Im} \psi_C | = \frac{ | (\mathbb{Z}/C\mathbb{Z}[i])^\times | }{ | \operatorname{Ker} \psi_C | } \] となります。\(| \operatorname{Im} \psi_C |\) こそが、我々が求めたかった 「\(\dfrac{a+bi}{a-bi}\) が取りうる値の種類数」\(g(C)\) です。

\(| (\mathbb{Z}/C\mathbb{Z}[i])^\times |\) は既知なので、\(h(C) \coloneqq |\operatorname{Ker} \psi_C|\) を求めます。 \[ \begin{aligned} & \, (x + yi) \in \operatorname{Ker} \psi_C \\ \iff & \, \begin{cases} (x + yi) \in (\mathbb{Z}/C\mathbb{Z}[i])^\times \\ \dfrac{x + yi}{x - yi} = 1 + 0i \end{cases} \\ \iff & \, \begin{cases} (x + yi) \in (\mathbb{Z}/C\mathbb{Z}[i])^\times \\ x + yi = x - yi \end{cases} \\ \iff & \, \begin{cases} \gcd(x^2 + y^2, C) = 1 \\ 2y \equiv 0 \pmod C \end{cases} \\ \iff & \, {}^\forall p \in \mathcal{P}( C ),\, \begin{cases} x^2 + y^2 \not\equiv 0 \pmod p \\ 2y \equiv 0 \pmod{p^{\nu _ p( C )}} \end{cases} \end{aligned} \] となります。 このことから、\(p \in \mathcal{P}(C)\) ごとに独立に \((x \bmod p^{\nu _ p(C)}, y \bmod p^{\nu _ p(C)})\) の取りうる値の種類数を求めることができて、 しかも中国剰余定理から、\((x, y)\) の取りうる値の種類数はそれらの積になります。 したがって、\(h(C)\) は \(C\) の乗法的関数となり、\(g(C) = \dfrac{f(C)}{h(C)}\) もまた乗法的関数になることが示されました。

素数 \(p\) と正整数 \(e\) を固定して、\(C = p^e\) に対して \((x, y)\) の取りうる値の種類数 \(h(p^e)\) を求めます。

  • \(p = 2\) のとき、 \(x^2 + y^2 \not\equiv 0 \pmod 2\) かつ \(y = 0, 2^{e-1}\) なので、\(x\) は奇数です。
    したがって、\((x, y)\) が取りうる値の種類数は \(2^e\) 個です。
  • \(p > 2\) のとき、 \(x^2 + y^2 \not\equiv 0 \pmod p\) かつ \(y \equiv 0 \pmod{p^e}\) なので、\(x\not\equiv0\pmod p\) です。
    したがって、\((x, y)\) が取りうる値の種類数は \(p^{e-1} (p-1)\) 個です。

これにより、 \[ g(p^e) = \frac{f(p^e)}{h(p^e)} = \begin{cases} \dfrac{2^{2e-1}}{2^e} & (p = 2) \\[10pt] \dfrac{(p^{e-1})^2 (p-1)^2}{p^{e-1} (p-1)} & (p \equiv 1 \pmod 4) \\[10pt] \dfrac{(p^{e-1})^2 (p^2-1)}{p^{e-1} (p-1)} & (p \equiv 3 \pmod 4) \end{cases} = \begin{cases} 2^{e-1} & (p = 2) \\ p^{e-1} (p-1) & (p \equiv 1 \pmod 4) \\ p^{e-1} (p+1) & (p \equiv 3 \pmod 4) \end{cases} \] となって、与式が示されます。

最後の結果は、\(\dfrac{a + bi}{a - bi}\) が取りうる値の種類数が \(C\) を高々素因数ごとに \(\dfrac{p+1}{p}\) 倍した程度にとどまることを示唆しています。 すなわち、この個数はほぼ \(O(C)\) である（より厳密には、\(O(C^{1 + \epsilon})\) である）のです。

さて、\(\bmod\) \(C\) で \[ (Q _ 0 + i R _ 0) \prod _ {j=1}^M (Q _ j + i R _ j)^{k _ j} (Q _ j - i R _ j)^{E _ j - k _ j} \quad\ldots(\dagger) \] が取りうる値の種類数の上界を求めます。
（この種類数が、上記のアルゴリズムで \(\mathrm{dp}\) および \(\mathrm{cnt}\) の添字 \((x, y)\) が取りうる添字ペアの個数に対応します。）
まず、素数 \(p\) と正整数 \(e\) を固定し、\(C = p^e\) であるときの上界を考えます。

2-1-4-(i). \(Q_j^2 + R_j^2\) がいずれも \(p\) と等しくないとき、 \((\dagger)\) が \(\bmod\) \(p^e\) で取りうる値の種類数が高々 \(g(p^e)\) 個であることを示します。 このとき、\(p\) と \(Q_j^2+R_j^2\) が相異なる素数であることから \(\gcd(Q_j^2+R_j^2, p^e)=1\)、すなわち \(Q_j \pm i R_j \in (\mathbb{Z}/p^e\mathbb{Z}[i])^\times\) なので、 \[ (\dagger) = \left( (Q _ 0 + i R _ 0) \prod _ {j=1}^M (Q _ j - i R _ j)^{E _ j} \right) \prod _ {j=1}^M \left( \frac{Q _ j + i R _ j}{Q _ j - i R _ j} \right)^{k _ j} \] となります。このうち前半は \(\mathbb{Z}/p^e\mathbb{Z}[i]\) の定数であり、後半は \((k_1, \ldots, k_M)\) にかかわらず \(\operatorname{Im} \psi_{p^e}\) に含まれます。 したがって、\((\dagger)\) が取りうる値の種類数は高々 \(|\operatorname{Im} \psi_{p^e}| \leq g(p^e)\) 個です。

2-1-4-(ii). \(Q_j^2 + R_j^2\) のいずれかが \(p\) と等しいとき、 \((\dagger)\) が \(\bmod\) \(p^e\) で取りうる値の種類数が高々 \(2 p^e\) 個であることを示します。 一般性を失わず \(Q_1^2 + R_1^2 = p\) とおきます。 \((\dagger)\) を \((x + yi)\) とおいたとき、 \[ \min(e, \nu_p(\gcd(x, y))) = \min(k_1, E_1 - k_1) \eqqcolon e^\prime \] がわかります。
直感的な説明 \((Q_1 + i R_1) (Q_1 - i R_1) = p\) から、

  • \((Q_1 + i R_1)^0 (Q_1 - i R_1)^{E_1 }\) は実部と虚部が \(1\) の倍数、
  • \((Q_1 + i R_1)^1 (Q_1 - i R_1)^{E_1 - 1}\) は実部と虚部が \(p\) の倍数、
  • \((Q_1 + i R_1)^2 (Q_1 - i R_1)^{E_1 - 2}\) は実部と虚部が \(p^2\) の倍数、
  • \(\quad\vdots\)
  • \((Q_1 + i R_1)^ e (Q_1 - i R_1)^{E_1 - e}\) は実部と虚部が \(p^e\) の倍数、すなわち \(0\)、
  • \((Q_1 + i R_1)^{e+1} (Q_1 - i R_1)^{E_1-e-1}\) は実部と虚部が \(p^e\) の倍数、すなわち \(0\)、
  • \(\quad\vdots\)
  • \((Q_1 + i R_1)^{E_1-e-1} (Q_1 - i R_1)^{e+1}\) は実部と虚部が \(p^e\) の倍数、すなわち \(0\)、
  • \((Q_1 + i R_1)^{E_1 - e} (Q_1 - i R_1)^ e \) は実部と虚部が \(p^e\) の倍数、すなわち \(0\)、
  • \(\quad\vdots\)
  • \((Q_1 + i R_1)^{E_1 - 2} (Q_1 - i R_1)^2\) は実部と虚部が \(p^2\) の倍数、
  • \((Q_1 + i R_1)^{E_1 - 1} (Q_1 - i R_1)^1\) は実部と虚部が \(p\) の倍数、
  • \((Q_1 + i R_1)^{E_1 } (Q_1 - i R_1)^0\) は実部と虚部が \(1\) の倍数、

となっています。

したがって、

  • \(e' < e\) のときは \((x/p^{e'}, y/p^{e'})\) に \(\bmod\) \(p^{e-e'}\) での自由度が残ります。
    しかし、(i) での議論から \((Q_2, R_2), \ldots, (Q_M, R_M)\) の寄与によって \((x,y)\) が取りうる値は \(g(p^{e - e'})\) 通りに限られます。
  • \(e' = e\) のときは \((x, y) = (0, 0)\) の \(1\) 通りに定まります。

ところで、\(Q_1^2 + R_1^2 = p\) より \(p \not\equiv 3 \pmod 4\) なので、\(g(p^{e - e'}) = p^{e - e^\prime - 1} (p-1)\) です。 以上を踏まえると、\((x, y)\) が取りうる値の種類数の上界は \[ 2 \sum_{e^\prime = 0}^{e - 1} p^{e-e^\prime-1} (p - 1) + 1 = 2 \frac{p^e - 1}{p - 1} \cdot (p - 1) + 1 \leq 2 p^e - 2 + 1 \leq 2 p^e \] となります。

したがって、一般の \(C\) に対しては \[ \begin{aligned} \, & (\dagger)\text{ が }\bmod C\text{ で取りうる値の種類数} \\ \leq \, & \prod _ {p \in \mathcal P( C )} (\dagger)\text{ が }\bmod p^{\nu _ p( C )} \text{ で取りうる値の種類数} \\ \leq \, & \prod _ {p \in \mathcal P( C )} 2 p^{\nu _ p( C )} \\ = \, & C \cdot \prod _ {p \in \mathcal P( C )} 2 \\ = \, & C \cdot (C\text{ の約数のうち平方因子を持たないものの個数}) \\ \leq \, & C \cdot (C\text{ の約数の個数}) \\ = \, & C \cdot d( C ) \end{aligned} \] が示されます。

よって、上記に示したアルゴリズム中の二次元配列 \(\mathrm{dp}\) および \(\mathrm{cnt}\) を連想配列に置き換えると、その要素数は高々 \(C \cdot d(C)\) となります。
また、\((\star)\) が最初と末尾の \(C \cdot d(C)\) 項を除くと周期 \(C \cdot d(C)\) 以下の周期性を持つことから、 \(\mathrm{cnt}\) は計算量 \(O(C \cdot d(C))\) で求められます。
したがって、これにより上記のアルゴリズムの計算量が \(O(M (C \cdot d(C))^2)\) となります。

\(d(n) = O(n^{(\log 2) / (\log \log n)})\) であることが知られており（参考）、 これは \(\operatorname{polylog}(n)\) よりは速く発散しますが、任意の \(\epsilon\) に対して \(n^{\epsilon}\) よりは遅く発散します。
これを用いると、全体の計算量は \[ O(M (C \cdot d( C ))^2 + P^{1.5}) = O(MC^{2 + 2 (\log 2) / (\log \log C)} + P^{1.5}) = O(MC^{2+\epsilon} + P^{1.5}) = o(MC^4 + P^{1.5}) \] と評価され、 元のアルゴリズムより高速であることがわかります。

（高速化を意識していなくても、二次元配列の代わりに連想配列やリストを用いる発想は非常に自然なものでしょう。 実際、提出されているコードのうち非常に高速に動作しているものの多くは \(\mathrm{dp}\) および \(\mathrm{cnt}\) （の一方あるいは両方）を二次元配列でなく連想配列やリストなどで管理しているようです。）

3. \(O(M C^2 \log \log C + P^{1.5})\) 解法

（注記：これも \(O(M C^{2+\epsilon} + P^{1.5})\) ではありますが、先程の解法より高速です。）

先程は（証明せず実装するだけなら）考察のいらない高速化を説明しましたが、 簡単な考察をすることで \(\mathrm{dp}\) および \(\mathrm{cnt}\) の取りうる添字 \((x, y)\) の種類数を \(O(C \sqrt{\log \log C})\) 個に抑えることができます。

3-1. 具体的な説明

まず自明な帰着として、\(g \coloneqq \gcd(A, B, C) > 1\) ならば \((A, B, C, N) \gets (A/g, B/g, C/g, N/g^2)\) のケースに帰着できます。
実装上は、各 \(p \in \mathcal{P}(g)\) に対して \(Q_j^2 + R_j^2 = p\) に対応する \(E_j\) から \(2 \nu_p(g)\) を引けば良いです。（引いた結果が負になる場合、答えは \(0\) です。）
（\(\mathcal P, \nu_p\) は 2-1. に定義されています。）

3-1-1. 次に、\((Q_j^2 + R_j^2) \mid C\) なる \((Q_j, R_j)\) が存在する場合を考え、これが存在しない場合に帰着することができることを示します。 そのような \((Q_j, R_j)\) を一つとり、\(p \coloneqq Q_j^2 + R_j^2\) とおきます。

2-1. の \((\dagger)\) をよく観察すると、\(k_j \geq 1\) かつ \(E_j - k_j \geq 1\) の場合 \((\dagger)\) の実部と虚部はともに \(p\) の倍数となります。
しかし、今 \(\gcd(A, B, C) = 1\) に帰着したので、そのような場合が答えに寄与することはありません。
したがって、\(k_j = 0, E_j\) に限定できます。

今、\(s = Q_j R_j^{-1}\) とします。（\(R_j^{-1}\) は \(\bmod\) \(p\) での逆元で、\(0 < R_j < \sqrt{p} < p\) より必ず取ることができます。）
また、ガウス整数から \(\mathbb Z/p\mathbb Z\) への写像 \(\varphi\) を \((x + yi) \mapsto x - ys\) で定義すると、これは群準同型をなすことが確認できます （\(s^2 \equiv -1 \pmod p\) に注意してください）。
すると、\(\varphi(Q_j + i R_j) = 0\) なので、任意のガウス整数 \(z\) に対して \(\varphi((Q_j + i R_j)z) = 0\) です。
したがって、\(\varphi(A + Bi) = 0\) ならば \(k_j = E_j\) が必要で、それ以外の場合は \(k_j = 0\) が必要になります。

これにより、\(k_j\) を一意に定めることができるため、\(Q_0 + i R_0\) を適切に置き換えることによって、\((Q_j, R_j)\) を削除することができます。
これを繰り返すことによって、\((Q_j^2 + R_j^2) \mid C\) なる \((Q_j, R_j)\) が存在しない場合に帰着することができます。

すると、2-1-4-(ii). のケースが排除できて、すべての素因数について 2-1-4-(i). の上界を適用できます。すなわち、 \[ \begin{aligned} \, & \text{2-1. の }(\dagger)\text{ が }\bmod C\text{ で取りうる値の種類数} \\ \leq \, & \prod _ {p \in \mathcal P( C )} \text{2-1. の }(\dagger)\text{ が }\bmod p^{\nu _ p( C )} \text{ で取りうる値の種類数} \\ \leq \, & \prod _ {p \in \mathcal P( C )} g(p^{\nu _ p( C )}) \\ \leq \, & \prod _ {\substack{p \in \mathcal P( C ) \\ p \equiv 3 \pmod 4}} p^{\nu _ p( C ) -1} (p + 1) \\ = \, & C \cdot \prod _ {\substack{p \in \mathcal P( C ) \\ p \equiv 3 \pmod 4}} \left(1 + \frac1p\right) \end{aligned} \] と上界を強くすることができます。

3-1-2. これが \(O(C \sqrt{\log\log C})\) であることを示します。

\[ f(n) \coloneqq \prod _ {\substack{p \in \mathcal P( n ) \\ p \equiv 3 \pmod 4}} \left(1 + \frac1p\right) \] とおきます。 テイラー展開および整数の逆二乗和が収束することから、 \[ \log f(n) = \sum _ {\substack{p \in \mathcal P( n ) \\ p \equiv 3 \pmod 4}} \log \left(1 + \frac1p\right) = \sum _ {\substack{p \in \mathcal P( n ) \\ p \equiv 3 \pmod 4}} \left(\frac1p + O\left(\frac1{p^2}\right) \right) = \sum _ {\substack{p \in \mathcal P( n ) \\ p \equiv 3 \pmod 4}} \frac1p + O(1) \] となります。\(n\) を固定したときに \(n' \leq n\) の範囲でこれを最大化するためには、 \(n'\) を \(4\) で割って \(3\) 余る素数を小さいものから順に掛けた積に取ればいいです。
このような積に含まれる最大の素数を \(p_{\max}\) と置くと、 算術級数の素数定理より \[ \sum _ {\substack{p \leq p _ {\max} \\ p\equiv3\pmod4}} \log p \sim \frac{p _ {\max}}{\varphi(4)} = \frac{p _ {\max}}{2} \] ですから、\(p_{\max} \sim 2 \log n\) です。（ただし、\(\varphi(n)\) はオイラーのトーシェント関数です。）
一方、定量的なディリクレの算術級数定理（ステートメントと証明 (Theorem 6.2, p.171)）より、 \[ \sum _ {\substack{p \leq p _ {\max} \\ p\equiv3\pmod4}} \frac1p = \frac1{\varphi(4)}\log\log p _ {\max} + O(1) = \frac12\log\log p _ {\max} + O(1) \] です。以上より、 \[ f(n) = \exp\left( \frac12\log\log p _ {\max} + O(1) \right) = \exp\left( \frac12\log\log(2 \log n) + O(1) \right) = O(\sqrt{\log\log n}) \] となって示されました。

この工夫により、全体の計算量を \(O(M C^2 \log \log C + P^{1.5})\) に削減できていることが、先ほどと同様の議論により示せます。

3-2. 実装上の補足

従来の解法と同様、\((\star)\) は最初と末尾を除くと周期性を持つことはすぐに分かります。
しかし実は、3-1. に示した工夫によって \((\star)\) は完全な周期性を持つようになりました。
なぜなら、（2-1-4-(i). の議論と同様の分解をすると） \[ (\star) = (Q _ j - i R _ j)^{E _ j} \left( \frac{Q _ j + i R _ j}{Q _ j - i R _ j} \right)^{k _ j} \] のうち前半は \(\mathbb{Z}/C \mathbb{Z}[i]\) の定数、後半は \(\operatorname{Im}\psi_C\) のある元の累乗であり、 \(\operatorname{Im}\psi_C\) が有限群をなしているからです。
（有限群の任意の元 \(g \in G\) に対して列 \((g, g^2, g^3, \ldots)\) を考えると鳩の巣原理により \(g^a = g^b\) なる正の整数 \(a, b\) \((|a - b| \leq |G|)\) が存在します。\(G\) が群であることから \((g^b)\) の逆元 \(g^{-b}\) が取れるのでこれを掛けると \(g^{a-b} = e\) （単位元）となることから、この列は \(|a-b|\) を周期にもちます。）

これにより、周期を元に寄与を計算する部分の実装を単純化することができます。

4. \(O(M C \log C \cdot \sqrt{\log \log C} + P^{1.5})\) 解法

現在の解法のボトルネックは \(\mathrm{dp}\) と \(\mathrm{cnt}\) のマージを愚直に行っている点です。 実は、このマージは畳み込みを活用すると、一回あたり \(O((\text{要素数}) \log(\text{要素数}))\) 時間で行うことができます。

4-1. 具体的な説明

3-1. に示した工夫によってこの問題は以下の問題に帰着されています。（2-1-4-(i). に示した表式と見比べてください。）

整数 \(A, B, C\)、非負整数列 \((E_1, \ldots, E_M)\)、 および \(z_0 \in \mathbb Z/C\mathbb Z[i]\) と \(z_1, \ldots, z_M \in \operatorname{Im}\psi_C (\subseteq \mathbb Z/C\mathbb Z[i])\) が与えられる。 非負整数列 \((k_0, \ldots, k_M)\) (\(k_0 < 4\), \(k_j \leq E_j\) (\(1 \leq j \leq M\))) であって、\(\mathbb Z/C\mathbb Z[i]\) 上の乗算について \[ i^{k_0} z_0 \prod _ {j=1}^M z _ j^{k _ j} = A + Bi \] を満たすものの個数を求めよ。 （\(Z/C\mathbb Z[i]\) は 2-1. に、\(\psi_C\) は 2-1-3. に定義されています。）

ここで、\(G \coloneqq \operatorname{Im}\psi_C\) は有限アーベル群（可換群）です。 一般に、有限アーベル群については以下の定理が知られています：

任意の空でない有限アーベル群 \(H\) に対して、正の整数 \(n_1, n_2, \ldots, n_m\) が存在して、 \[ H \cong (\mathbb Z/n_1\mathbb Z) \times (\mathbb Z/n_2\mathbb Z) \times \cdots \times (\mathbb Z/n_m\mathbb Z) \] が成り立つ。 ただし \((\mathbb Z/n_j\mathbb Z)\) は加法に関する群である。

直感的には、可換な群は「\({}\bmod n_1, {}\bmod n_2, \ldots, {}\bmod n_m\) の整数列」とみなすことができるということです。 より具体的に言うと、ある \(G\) の元の列 \((g_1, \ldots, g_m)\) を取ることができて、 任意の \(z \in G\) に対して非負整数 \(e_j < |g_j|\) \((j = 1, \ldots, m)\) が一意に存在して \[ z = g_1^{e_1} g_2^{e_2} \cdots g_m^{e_m} \quad \ldots(\clubsuit)\] と書くことができます。 （\(|g_j|\) は \(G\) 上での \(g_j\) の位数、すなわち \(g_j^e\) が単位元に等しくなる最小の正の整数 \(e\) です。)

いま、 \[ (g_1^{e_1} g_2^{e_2} \cdots g_m^{e_m}) (g_1^{e_1^\prime} g_2^{e_2^\prime} \cdots g_m^{e_m^\prime}) = (g_1^{(e_1 + e_1^\prime) \bmod n_1} g_2^{(e_2 + e_2^\prime)\bmod n_2} \cdots g_m^{(e_m + e_m^\prime)\bmod n_m}) \] なる関係に着目すると、これは \(R\)-上 \(m\) 変数多項式であって \(x_j^{n_j} = 1\) となるように剰余を取ったもの \[ P( R ) \coloneqq R[x_1, \ldots, x_m] / (1 - x_1^{n_1}) (1 - x_2^{n_2}) \cdots (1 - x_m^{n_m}) \] 上の積に対応させることができます。 したがって、\(\mathrm{dp}\) や \(\mathrm{cnt}\) もこの多項式の剰余環 \(P\) の係数にエンコードすることができ、 \(\mathrm{dp'}\) は単に対応する多項式の積として得ることができます。 （これを 巡回畳み込み と言います。）

\(n \coloneqq n_1 n_2 \cdots n_m\) とすると、このような \(m\) 変数の巡回畳み込みは 環 \(R\) に次数 \(n_1, n_2, \ldots, n_m\)、および \(2^{\lceil \log_2 \max_i n_i \rceil + 1}\) の原始根が存在するならば \(O(n \log n)\) 時間で行えることが知られています （参考：37zigen さんの記事、 Library Checker）。 今回は \(\bmod\) \(998244353\) での計算が要求されているので、\(kn + 1\) の形で表されるいくつかの素数の集合であって、総積が \(998244352^2\) を超えるものを一つとり、 それらを法とする有限体上の多項式の畳み込みを行ったのち、Garner のアルゴリズムによって正しい値を復元すると、所望の値が得られます。

残る問題は、\(G\) の巡回群への分解を行う方法、すなわち

  • \((g_1, \ldots, g_m)\) を得るアルゴリズム
  • 任意の \(g \in G\) に対して \((\clubsuit)\) を満たす整数列 \((e_1, \ldots, e_m)\) を得るアルゴリズム

です。特に \(G\) の構造の解析や元の列挙は自明なタスクではありません。\(G\) の要素数 \(|G|\) すら知りません。 いい方法はないでしょうか？

4-1-1. 今回は前計算に \(\tilde O(|G|)\) 時間程度かけても問題ないため、以下のようなアルゴリズムが使用できます。

(入力) \(G\) の単位元 \(z_0\) および \(G\) の元 \(z_1, \ldots, z_k\)

(出力) \(G'\) の元の列 \((g_1, \ldots, g_m)\) であって、任意の \(z \in G'\) に対して \((\clubsuit)\) の表現が一意となるもの （ただし \(G'\) は \(z_1, \ldots, z_k\) で生成される \(G\) の部分群）

(アルゴリズム)

  1. \(G'\) の元をすべて列挙する。
    • \(G'\) の元からなる集合 generated をはじめ \(\{z_0\}\) で初期化する。
    • \(j = 1, \ldots, k\) の順に、以下を行う。
      • \(z_j \notin\) generated なら、以下を行う。
        • （現時点での）各 \(z \in\) generated に対して、以下を行う。
          • \(z' = z_j^1 z, z_j^2 z, \ldots\) を列挙する。それぞれに対して、
            • \(z'\) が generated に含まれないなら、generated に追加する。
            • 含まれるなら、この列挙ループを終了する。
    • これにより、generated は \(G'\) の全ての元を含む。
  2. \(|G'|\) およびその素因数分解を求める。
    • generated のサイズを \(|G'|\) と認識する。
    • \(|G'|\) 素因数分解し \(p_1^{e_1} p_2^{e_2} \cdots p_l^{e_l}\) とおく。
  3. \(G'\) の各元の位数を求める。
    • 整数から \(G'\) の元のリストへの連想配列 of_order をはじめ空で初期化する。
    • 各 \(z \in\) generated に対して、以下を行う。
      • \(|G'|\) の約数 \(d\) を小さい方から列挙する。
        • \(z^d = z_0\) ならば、of_order\([d]\) に \(z\) を追加して、この列挙ループを終了する。
    • これにより、全ての \(z \in G\) は of_order のいずれかのリストに含まれ、そのキーは \(|z|\) に等しくなる。
  4. \((g_1, \ldots, g_m)\) を求める。
    • \(G'\) の元からなるリスト generators をはじめ空リストで初期化する。
    • \(j = 1, \ldots, l\) の順に、以下を行う。
      • \(G\) の元からなる集合 generated をはじめ \(\{z_0\}\) で初期化する。
      • \(e = e_j, e_j - 1, \ldots, 1\) の順に、以下を行う。
        • 各 \(z \in \) of_order\([p_j^e]\) に対して、以下を行う。
          • generated に \(z^{p_j^{e-1}}\) が含まれてなければ、以下を行う。
            • generators に \(z\) を追加する。
            • (現時点での) 各 \(x \in\) generated に対して、以下を行う。
              • generated に \(xz, xz^2, \ldots, xz^{p_j^e - 1}\) を追加する。
    • これにより、generators は求める \((g_1, \ldots, g_m)\) としての条件を満たすので、これを出力する。

(計算量) 全体としての計算量は \(O(|G| \log |G| + k)\) 程度です。特にステップ 1. については非自明ですが、以下のように解析されます。

  • \(j\) がループする回数はちょうど \(k\) 回である。
  • \(j\) についてのループが終了したとき、generated は \(z_1, \ldots, z_j\) で生成される \(G\) の部分群に等しい。
  • \(z_j \notin\) generated が発生するとき、部分群 generated は真に大きくなる。この処理中で、
    • 「含まれないなら、」が発生する合計回数は generated の新しいサイズと古いサイズの差に等しい。
      • したがって、そのアルゴリズム全体にわたる合計は \(|G|\) に等しい。
    • 「含まれるなら、」が発生する合計回数は generated の古いサイズに等しい。
      • generated のサイズは、 ラグランジュの定理よりちょうど整数倍、 特に \(2\) 倍以上になる。
      • したがって、この古いサイズの合計は \(\dfrac{|G|}{2} + \dfrac{|G|}{4} + \dfrac{|G|}{8} + \cdots \leq |G|\) 以下である。
  • したがって、このステップの計算量は \(O(|G| + k)\) である。

\((g_1, \ldots, g_m)\) が得られたら、それらのすべての指数を列挙して乗算することで、 \((\clubsuit)\) を満たす整数列 \((e_1, \ldots, e_m)\) との対応付けも行うことができます。

要素数が \(O(C \sqrt{\log \log C})\) であったことを踏まえれば、これにより \(O(M C \log C \cdot \sqrt{\log \log C} + P^{1.5})\) 解法が得られます。

5. 参考文献

Allan, Adam A., et al. “Classification of the group of units in the Gaussian integers modulo n.” Pi Mu Epsilon Journal 12.9 (2008): 513-519.

6. さらなる高速化に対する疑問

今、問題は以下に帰着されました。

非負整数 \(k\)、正整数 \(n\)、\(2\) 以上の整数からなる列 \((n_1, \ldots, n_k)\) (ただし \(\prod_{j=1}^k n_j = n\)）、\((m\times k)\)-非負整数行列 \((a_{ij})\) （ただし \(a_{ij} < n_j\)）、非負整数列 \((e_1, \ldots, e_m)\) が与えられる。 \[ \prod _ {i=1}^m \sum _ {e=0}^{e _ i} \prod _ {j=1}^k x_j^{a _ {ij}e} \in \mathbb F _ {998244353} [x _ 1, \ldots, x _ k] / (1 - x _ 1^{n _ 1}) (1 - x _ 2^{n _ 2}) \cdots (1 - x _ k^{n _ k}) \] のうち、指定された \(O(1)\) 個の係数を求めよ。

この問題が \(o(m n \log n)\) 時間で解ければ、元の問題もさらに高速に解くことができるのですが、可能でしょうか？
筆者はこれに対する解法の有無を知りません。

投稿日時:
最終更新: