ゼータ変換・素数ゼータ変換について

ゼータ変換とは

ゼータ変換とは、

$n$ 次元累積和のことです。

$2$ 次元の場合を考えてみましょう。

$n\times m$ の数列

$A$ に対し、

$$B[x][y]=\sum _{i=0}^x\sum _{j=0}^{y}A[i][j]$$

であるような

$n\times m$ の数列

$B$ を計算することを考えます。

1	0	2
0	3	0
4	0	5

$⟹$

1	1	3
1	4	6
5	8	15

これは、

$$C[x][y]=\sum _{j=0}^{y}A[x][j]$$

を定義すると

$$B[x][y]=\sum _{i=0}^x(\sum _{j=0}^{y}A[i][j])=\sum _{i=0}^{x}C[i][y]$$

なので、片方の次元について累積和をした後、もう片方の次元について累積和をすると求まります。

1	0	2
0	3	0
4	0	5

$⟹$
横

1	1	3
0	3	3
4	4	9

$⟹$
縦

1	1	3
1	4	6
5	8	15

$n$ 次元累積和（ゼータ変換）がしたいときには、ある次元について累積和を取ることを

$n$ 回別々に行えばよいです。

注意点

複数の次元で同時に累積和を取ろうとしてはいけません。

# 最初の次元についての累積和
for i in range(N-1):
    for j in range(M):
        B[i+1][j] += B[i][j]

# 2 つ目の次元についての累積和
for i in range(N):
    for j in range(M-1):
        B[i][j+1] += B[i][j]

は正しいですが、

for i in range(N):
    for j in range(M):
        # 最初の次元についての累積和
        if i + 1 < N:
            B[i+1][j] += B[i][j]
        # 2 つ目の次元についての累積和
        if j + 1 < M:
            B[i][j+1] += B[i][j]

は WA です。(たまにやる)

素数ゼータ変換とは

素数の次元でゼータ変換をやります。

$A$ を、

$2$ べきの添字のみで定義される配列とします。

$A[1]$	$A[2]$	$A[4]$	$A[8]$	$A[16]$	$A[32]$	$A[64]$
2	0	3	1	0	5	0

これに対し、

$$B[x]=\sum _{y\text{ は }x\text{ の約数}}A[y]$$

で定義される配列

$B$ を求めることを考えます。

$B[1]$	$B[2]$	$B[4]$	$B[8]$	$B[16]$	$B[32]$	$B[64]$
2	2	5	6	6	11	11

これは、ただの累積和ですね。

---

$A$ を、(

$2$ べき)

$\times$ (

$3$ べき) の添字のみで定義される配列とします。

$A[1]$	$A[2]$	$A[3]$	$A[4]$	$A[6]$	$A[8]$	$A[9]$	$A[12]$
1	-1	4	-2	-4	-3	2	3

これを

$2$ 次元に直してみます。

$A$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	1	-1	-2	-3
$3^1$	4	-4	3
$3^2$	2

これに対し、

$$B[x]=\sum _{y\text{ は }x\text{ の約数}}A[y]$$

で定義される配列

$B$ を求めることを考えます。

$B$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	1	0	-2	-5
$3^1$	5	0	1
$3^2$	7

$B[1]$	$B[2]$	$B[3]$	$B[4]$	$B[6]$	$B[8]$	$B[9]$	$B[12]$
1	0	5	-2	0	-5	7	1

これは、

$2$ べきの方向に累積和を取ったあと、

$3$ べきの方向に累積和を取ることで求められます。

$A$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	1	-1	-2	-3
$3^1$	4	-4	3
$3^2$	2

$⟹$
横

$B$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	1	0	-2	-5
$3^1$	4	0	3
$3^2$	2

$⟹$
縦

$B$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	1	0	-2	-5
$3^1$	5	0	1
$3^2$	7

---

$A$ を、(

$2$ べき)

$\times$ (

$3$ べき)

$\times$ (

$5$ べき)

$\times \cdots$ の添字のみで定義される配列とします。

$A[1]$	$A[2]$	$A[3]$	$A[2^2]$	$A[5]$	$A[2\cdot 3]$	$A[7]$	$A[2^3]$	$A[3^2]$	$A[2\cdot 5]$	$A[11]$	$A[2^2\cdot 3]$
1	0	1	0	0	-1	1	1	0	-1	1	0

これに対し、

$$B[x]=\sum _{y\text{ は }x\text{ の約数}}A[y]$$

で定義される配列

$B$ を求めることを考えます。

これは、

$2$ べきの方向に累積和を取り、

$3$ べきの方向に累積和を取り、

$5$ べきの方向に累積和を取り、

$7$ べきの方向に累積和を取り、

$11$ べきの方向に累積和を取り…

とすることで求められます。

まとめ

長さ

$N$ の数列

$A[1],A[2],\dots ,A[N]$ に対し、

$$B[x]=\sum _{y\text{ は }x\text{ の約数}}A[y]$$

で定義される長さ

$N$ の数列

$B[1],B[2],\dots ,B[N]$ を求めるには、

1. $N$ 以下の素数を列挙する
2. $N$ 以下の各素数
   $p$ について:
   • $p$ べき方向の累積和を取る

計算量

• 1. はエラトステネスの篩で
  $O(N\log \log N)$ 時間 (線形でもできる)
• 2. はエラトステネスの篩と同様の計算量解析で、素数の逆数和なので
  $O(N\log \log N)$ 時間である。
• ややランダムアクセス気味なのでキャッシュに載りにくく、定数倍はまあまあある。

実装例

# primes: N 以下の素数のリスト (適当な篩で用意する)
for p in primes:
    for i in range(1, N // p + 1):
        B[i * p] += B[i]

---

倍数添字の和を求めたいとき

$A$ を、(

$2$ べき)

$\times$ (

$3$ べき) の添字のみで定義される配列とします。

$A[1]$	$A[2]$	$A[3]$	$A[4]$	$A[6]$	$A[8]$	$A[9]$	$A[12]$
1	-1	4	-2	-4	-3	2	3

これに対し、

$$B[x]=\sum _{y\text{ は }x\text{ の倍数}}A[y]$$

で定義される配列

$B$ を求めることを考えます。

逆向きに累積和をすれば良いです。

$A$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	1	-1	-2	-3
$3^1$	4	-4	3
$3^2$	2

$⟹$
横

$B$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	-5	-6	-5	-3
$3^1$	3	-1	3
$3^2$	2

$⟹$
縦

$B$	$2^0$	$2^1$	$2^2$	$2^3$
$3^0$	0	-7	-2	-3
$3^1$	5	-1	3
$3^2$	2

これでちゃんと倍数添字の和が求まっているか確認しておきましょう。

ABC393 E - GCD of Subset を解く

$x$ の倍数が

$K$ 個以上存在すれば GCD として

$x$ の倍数が作れるので、

cnt[x] = sum(a % x == 0 for a in A)

で定義される数列 cnt を求めます。これは、

for a in A:
    cnt[a] += 1

をした後、約数方向にゼータ変換

for p in primes:
    for i in range(MX // p, 0, -1):
        cnt[i] += cnt[i * p]

すれば良いです。

さて問題に戻ると、

$A_i$ を含む時点で GCD は

$A_i$ の約数に限られます。

$A_i$ の約数

$x$ について、

$\text{cnt}[x]<K$ ならば GCD を

$x$ の倍数にすることはできないし、

$\text{cnt}[x]\ge K$ ならば GCD を

$x$ の倍数にすることができます。

したがって、

$$max\{x\in \text{divisor}(A_i)\mid \text{cnt}[x]\ge K\}$$

が答えになります。これを求めるためには、

for i in range(1, MX + 1):
    if cnt[i] >= K:
        ans[i] = i

をした後、倍数方向にゼータ変換（今度は累積和ではなく、累積 max）をすれば良いです。

for p in primes:
    for i in range(1, MX // p + 1):
        if ans[i * p] < ans[i]:
            ans[i * p] = ans[i]

計算量は、

$O(A_{max}\log \log A_{max}+N)$ 時間です。

解答

https://atcoder.jp/contests/abc393/submissions/62825909

import sys
input = sys.stdin.readline

N, K = map(int, input().split())
A = list(map(int, input().split()))

MX = max(A)
# 線形篩
primes = []
factor = [1] * (MX + 1)
for i in range(2, MX + 1):
    if factor[i] == 1:
        primes.append(i)
        factor[i] = i
    for p in primes:
        if i * p > MX or p > factor[i]:
            break
        factor[i * p] = p

# cnt[x] = sum(a % x == 0 for a in A)
cnt = [0] * (MX + 1)
for a in A:
    cnt[a] += 1
for p in primes:
    for i in range(MX // p, 0, -1):
        cnt[i] += cnt[i * p]

# ans[x] = max(d for d in divisor(x) if cnt[d] >= K)
ans = [0] * (MX + 1)
for i in range(1, MX + 1):
    if cnt[i] >= K:
        ans[i] = i
for p in primes:
    for i in range(1, MX // p + 1):
        if ans[i * p] < ans[i]:
            ans[i * p] = ans[i]

for a in A:
    print(ans[a])