AtCoder Beginner Contest 349 F問題メモ

F - Subsequence LCM

問題

長さ $N$ の数列 $A=(A_1,...,A_N)$ がある
ここから1つ以上の要素を選んで、その最大公約数がちょうど $M$ になるようにしたい
選び方の個数を $\mod{998244353}$ で求めよ
- $A$ には同じ値も含まれうるが、異なる位置から採用したものは別物として数える
$1 \le N \le 2 \times 10^5$
$1 \le M \le 10^{16}$

解法1

最大公約数を考えるには、まず $M$ を素因数分解する。
$M \le 10^{16}$ なので $O(\sqrt{M})$ の素因数分解はやや厳しく見えるが、演算が単純なのでできる。
不安なら確率的になるが高速なポラード・ロー法を使うこともできる。

たとえば $M=1176=2^{3}3^{1}7^{2}$ とできる場合、

$M$ を割り切ることができない $A_i$ は採用できないので、はじめから除いて考えてよい
$2^{3}$ を素因数に含む $A_i$ を最低1つ含む必要がある
$3^{1}$ を素因数に含む $A_i$ を最低1つ含む必要がある
$7^{2}$ を素因数に含む $A_i$ を最低1つ含む必要がある

これさえ満たせば、最小公倍数が $M$ になる。
すると、少し特殊なナップサック問題みたいに言い換えられる。

まず、 $M$ の素因数の種類数を $K$ （今回は $3$ ）として、各 $A_i$ をもとに長さ $K$ の bitflag $B_i$ を作る。

$A_i$ が $2^{3}$ を素因数に含むとき、 $B_i$ の1の位に1を立てる
$A_i$ が $3^{1}$ を素因数に含むとき、 $B_i$ の10の位に1を立てる
$A_i$ が $7^{2}$ を素因数に含むとき、 $B_i$ の100の位に1を立てる

すると、「 $B_i$ からいくつかを選んで、その bitwise-or が長さ $K$ の 111…1 となるような選び方の個数」を数えればよいことになる。

このとき、 $2 \times 3 \times 5 \times 7 \times ...$ のように素数を小さい方からかけると $14$ 個で $10^{16}$ を超えるので、素因数の種類数は最大でも $K \le 13$ となる。

$2^{13}=8192$ なので、普通にナップサックDPすると $O(N 2^K) \simeq 1.6 \times 10^9$ となりダメだが、
種類数が少ないことを利用して $B_i$ が同じものはまとめて遷移させると $O(4^K) \simeq 6.7 \times 10^7$ で求められる。

Python3

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

from collections import Counter
 
# Pollard's rho algorithm
# https://qiita.com/t_fuki/items/7cd50de54d3c5d063b4a
 
def gcd(a, b):
    while a:
        a, b = b % a, a
    return b
 
 
def is_prime(n):
    if n == 2:
        return 1
    if n == 1 or n % 2 == 0:
        return 0
 
    m = n - 1
    lsb = m & -m
    s = lsb.bit_length() - 1
    d = m // lsb
 
    test_numbers = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]
 
    for a in test_numbers:
        if a == n:
            continue
        x = pow(a, d, n)
        r = 0
        if x == 1:
            continue
        while x != m:
            x = pow(x, 2, n)
            r += 1
            if x == 1 or r == s:
                return 0
    return 1
 
 
def find_prime_factor(n):
    if n % 2 == 0:
        return 2
 
    m = int(n ** 0.125) + 1
 
    for c in range(1, n):
        f = lambda a: (pow(a, 2, n) + c) % n
        y = 0
        g = q = r = 1
        k = 0
        while g == 1:
            x = y
            while k < 3 * r // 4:
                y = f(y)
                k += 1
            while k < r and g == 1:
                ys = y
                for _ in range(min(m, r - k)):
                    y = f(y)
                    q = q * abs(x - y) % n
                g = gcd(q, n)
                k += m
            k = r
            r *= 2
        if g == n:
            g = 1
            y = ys
            while g == 1:
                y = f(y)
                g = gcd(abs(x - y), n)
        if g == n:
            continue
        if is_prime(g):
            return g
        elif is_prime(n // g):
            return n // g
        else:
            return find_prime_factor(g)
    print('Hello', n)
 
 
def factorize(n):
    res = {}
    while not is_prime(n) and n > 1:  # nが合成数である間nの素因数の探索を繰り返す
        p = find_prime_factor(n)
        s = 0
        while n % p == 0:  # nが素因数pで割れる間割り続け、出力に追加
            n //= p
            s += 1
        res[p] = s
    if n > 1:  # n>1であればnは素数なので出力に追加
        res[n] = 1
    return res
 
 
n, m = map(int, input().split())
aaa = list(map(int, input().split()))
MOD = 998244353
 
if m == 1:
    ans = pow(2, aaa.count(1), MOD) - 1
    print(ans)
    exit()
 
fct = factorize(m)
s = len(fct)
mask = (1 << s) - 1
fct = [(i, pow(p, q)) for i, (p, q) in enumerate(fct.items())]
 
bbb = Counter()
for a in aaa:
    if m % a != 0:
        continue
    tmp = 0
    for i, pq in fct:
        if a % pq == 0:
            tmp |= 1 << i
    bbb[tmp] += 1
 
init = pow(2, bbb[0], MOD)
del bbb[0]
 
ss = 1 << s
dp = [0] * ss
dp[0] = init
 
for bit, cnt in bbb.items():
    pat = pow(2, cnt, MOD) - 1
    for i in range(ss - 1, -1, -1):
        dp[i | bit] += dp[i] * pat
        dp[i | bit] %= MOD
 
print(dp[mask])

解法2

bitwise-or のところまでは解法1と同じだが、ゼータ変換・メビウス変換を使えばナップサックより高速に求められる。

Bi ごとに、個数をカウントした配列
  000  001  010  011  100  101  110  111
[   3    2    0    1    2    3    0    0 ]

↓ゼータ変換: 
    たとえば 110 の値は、Bi=000,010,100,110 のように110の部分集合であるBiの個数の和

  000  001  010  011  100  101  110  111
[   3    5    3    6    5   10    5   11 ]

↓2の冪乗にする:
    たとえば 110 の値は、bitwise-orが110の部分集合となるようなBiの選び方の個数

  000  001  010  011  100  101  110  111
[   8   32    8   64   32 1024   32 2048 ]

↓メビウス変換:
    たとえば 110 の値は、bitwise-orが "ちょうど" 110となるようなBiの選び方の個数

  000  001  010  011  100  101  110  111
[   8   24    0   32   24  968    0  992 ]

よって、この例の答えは 992 となる。