結果
| 問題 |
No.2959 Dolls' Tea Party
|
| コンテスト | |
| ユーザー |
 qwewe
|
| 提出日時 | 2025-05-14 13:18:42 |
| 言語 | PyPy3 (7.3.15) |
| 結果 |
TLE
|
| 実行時間 | - |
| コード長 | 11,176 bytes |
| コンパイル時間 | 171 ms |
| コンパイル使用メモリ | 82,740 KB |
| 実行使用メモリ | 173,800 KB |
| 最終ジャッジ日時 | 2025-05-14 13:19:38 |
| 合計ジャッジ時間 | 5,534 ms |
|
ジャッジサーバーID （参考情報） |
judge4 / judge2 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 4 |
| other | AC * 2 TLE * 1 -- * 30 |
ソースコード
import sys
# Increase recursion depth for deep polynomial operations if necessary
# sys.setrecursionlimit(2000)
def solve():
N, K = map(int, sys.stdin.readline().split())
A = list(map(int, sys.stdin.readline().split()))
MOD = 998244353
# Precompute factorials, inverse factorials, and modular inverses
# Need values up to K+1 because integration involves division by k+1 where k can be up to K
MAX_K = K + 2
fact = [1] * MAX_K
inv_fact = [1] * MAX_K
inv = [1] * MAX_K
for i in range(1, MAX_K):
fact[i] = (fact[i-1] * i) % MOD
# Compute inverse factorial of MAX_K-1 using Fermat's Little Theorem
inv_fact[MAX_K-1] = pow(fact[MAX_K-1], MOD - 2, MOD)
# Compute inverse factorials downwards
for i in range(MAX_K - 2, -1, -1):
inv_fact[i] = (inv_fact[i+1] * (i+1)) % MOD
# Compute modular inverses using inverse factorials
# inv[i] = i^{-1} mod MOD = (i-1)! * (i!)^{-1} mod MOD
for i in range(1, MAX_K):
inv[i] = (inv_fact[i] * fact[i-1]) % MOD
# Technically inv[0] is undefined, set to 0 for safety though not used meaningfully
inv[0] = 0
# NTT implementation
G = 3 # Primitive root for MOD = 998244353
def ntt(a, is_inverse):
""" Performs Number Theoretic Transform (NTT) in place. """
n = len(a)
# Bit-reversal permutation
j = 0
for i in range(1, n):
rev = n >> 1
while j >= rev:
j -= rev
rev >>= 1
j += rev
if i < j:
a[i], a[j] = a[j], a[i]
# Cooley-Tukey algorithm
k = 1
while k < n:
# Twiddle factor step
w_step = pow(G, (MOD - 1) // (2 * k), MOD)
if is_inverse:
# Use inverse root for inverse NTT
w_step = pow(w_step, MOD - 2, MOD)
for i in range(0, n, 2 * k):
w = 1 # Current twiddle factor omega^j
for j in range(i, i + k):
u = a[j]
v = (a[j + k] * w) % MOD
a[j] = (u + v) % MOD
a[j + k] = (u - v + MOD) % MOD
w = (w * w_step) % MOD # Update twiddle factor
k *= 2
if is_inverse:
# Divide by n for inverse transform
n_inv = pow(n, MOD - 2, MOD)
for i in range(n):
a[i] = (a[i] * n_inv) % MOD
return a
def poly_mul(p1, p2, target_len):
""" Multiplies two polynomials p1 and p2 using NTT, returns result truncated to target_len. """
len1 = len(p1)
len2 = len(p2)
if len1 == 0 or len2 == 0: return [0] * target_len
# Determine required NTT length (power of 2 >= final degree + 1)
n = 1
while n < len1 + len2 - 1:
n <<= 1
# Pad polynomials to length n
p1_ntt = p1 + [0] * (n - len1)
p2_ntt = p2 + [0] * (n - len2)
# Perform NTT on both polynomials
ntt(p1_ntt, False)
ntt(p2_ntt, False)
# Pointwise multiplication in frequency domain
res_ntt = [(p1_ntt[i] * p2_ntt[i]) % MOD for i in range(n)]
# Inverse NTT to get result in coefficient domain
ntt(res_ntt, True)
# Return result truncated to target_len
return res_ntt[:target_len]
def poly_inv(p, target_len):
""" Computes modular inverse of polynomial p modulo x^target_len using Newton-Raphson. """
if target_len == 0: return []
# Pad p to target_len to avoid index errors and ensure consistency
p_eff = p[:target_len] + [0]*(target_len - len(p))
if target_len == 1:
# Base case: inverse of constant term
if not p_eff or p_eff[0] == 0:
raise ValueError("Polynomial has no inverse (constant term is 0)")
return [pow(p_eff[0], MOD - 2, MOD)]
# Recursive step: find inverse mod x^k first
k = (target_len + 1) // 2
inv_k = poly_inv(p, k) # Recursive call with original p, length k
# Determine NTT length
n = 1
while n < target_len + k -1: # Needs length at least target_len for result, and involves P up to target_len
n <<= 1
# Prepare P and Inv_k for NTT
p_ntt = p_eff[:target_len] + [0] * (n - target_len)
inv_k_ntt = inv_k + [0] * (n - len(inv_k))
ntt(p_ntt, False)
ntt(inv_k_ntt, False)
# Compute P * Inv_k mod x^target_len
pk_ik_ntt = [(p_ntt[i] * inv_k_ntt[i]) % MOD for i in range(n)]
ntt(pk_ik_ntt, True)
pk_ik = pk_ik_ntt[:target_len]
# Newton-Raphson update step: Calculate 2 - P * Inv_k mod x^target_len
res_poly = [0] * target_len
res_poly[0] = 2
for i in range(target_len):
res_poly[i] = (res_poly[i] - pk_ik[i] + MOD) % MOD
# Transform (2 - P*Inv_k) to frequency domain
res_poly_ntt = res_poly + [0] * (n - target_len)
ntt(res_poly_ntt, False)
# Multiply Inv_k * (2 - P * Inv_k) in frequency domain
final_res_ntt = [(inv_k_ntt[i] * res_poly_ntt[i]) % MOD for i in range(n)]
# Transform back to coefficient domain
ntt(final_res_ntt, True)
return final_res_ntt[:target_len]
def poly_deriv(p):
""" Computes the derivative of polynomial p. """
if len(p) <= 1: return []
res = [0] * (len(p) - 1)
for i in range(1, len(p)):
res[i-1] = (p[i] * i) % MOD
return res
def poly_integr(p):
""" Computes the integral of polynomial p (constant term 0). """
res = [0] * (len(p) + 1)
for i in range(len(p)):
# Needs inv[i+1] precomputed
res[i+1] = (p[i] * inv[i+1]) % MOD
return res
def poly_log(p, target_len):
""" Computes logarithm ln(P(x)) mod x^target_len. Requires P(0)=1. """
# Pad p to target_len for internal calculations
p_eff = p[:target_len] + [0]*(target_len - len(p))
# Check prerequisite P(0)=1
if not p_eff or p_eff[0] != 1:
raise ValueError(f"Logarithm undefined for polynomial not starting with 1. P[0]={p_eff[0] if p_eff else 'empty'}")
if target_len == 1: return [0] # log(1) = 0
# Compute P'(x)
p_deriv = poly_deriv(p_eff)
# Compute P(x)^{-1} mod x^target_len
p_inv = poly_inv(p_eff, target_len)
# Compute P'(x) * P(x)^{-1} mod x^target_len
prod = poly_mul(p_deriv, p_inv, target_len)
# Integrate the product. Integration needs terms up to x^(target_len-2) to get result up to x^(target_len-1).
# Integral of Prod[k] * x^k is Prod[k]/(k+1) * x^(k+1).
res = poly_integr(prod[:target_len-1])
return res[:target_len]
def poly_exp(q, target_len):
""" Computes exponentiation exp(Q(x)) mod x^target_len. Requires Q(0)=0. """
# Pad q to target_len
q_eff = q[:target_len] + [0]*(target_len - len(q))
# Check prerequisite Q(0)=0
if not q_eff or q_eff[0] != 0:
raise ValueError(f"Exponentiation undefined for polynomial not starting with 0. Q[0]={q_eff[0] if q_eff else 'empty'}")
if target_len == 1: return [1] # exp(0) = 1
# Recursive step: find exp(Q) mod x^k first
k = (target_len + 1) // 2
exp_k = poly_exp(q_eff, k) # Recursive call with q_eff, length k
# Compute ln(exp(Q)_k) mod x^target_len. exp_k has length k.
log_exp_k = poly_log(exp_k, target_len) # Pass exp_k directly, poly_log handles padding
# Newton-Raphson update: Compute 1 + Q - ln(Exp_k) mod x^target_len
res_poly = [0] * target_len
res_poly[0] = 1
# Add Q
for i in range(target_len):
res_poly[i] = (res_poly[i] + q_eff[i]) % MOD
# Subtract ln(Exp_k)
for i in range(target_len):
res_poly[i] = (res_poly[i] - log_exp_k[i] + MOD) % MOD
# Multiply Exp_k * (1 + Q - ln(Exp_k)) mod x^target_len
final_res = poly_mul(exp_k, res_poly, target_len)
return final_res[:target_len]
# Euler totient function calculation
phi_map = {}
def get_phi(m):
""" Computes Euler's totient function phi(m), uses caching. """
if m in phi_map: return phi_map[m]
if m == 0: return 0
if m == 1: phi_map[1] = 1; return 1
res = m
p = 2
temp_m = m
while p * p <= temp_m:
if temp_m % p == 0:
res = res // p * (p - 1)
while temp_m % p == 0:
temp_m //= p
p += 1
if temp_m > 1: # If temp_m is still > 1, it must be a prime factor
res = res // temp_m * (temp_m - 1)
phi_map[m] = res
return res
# Find all divisors of K
divisors = []
for i in range(1, int(K**0.5) + 1):
if K % i == 0:
divisors.append(i)
if i*i != K:
divisors.append(K//i)
divisors.sort()
# Main logic using Burnside's Lemma
total_S = 0 # Sum of |X^g| terms
for d in divisors:
L = K // d # Cycle length
# Group doll types by B_i = floor(A_i / L) value
counts = {} # counts[b] = number of types i with B_i = b
for x in A:
Bi = x // L
if Bi > 0: # Only types that can form at least one cycle matter
counts[Bi] = counts.get(Bi, 0) + 1
# Calculate Sum_{b} C_b * ln(P_b(x)) where P_b(x) = sum_{k=0..b} x^k/k!
sum_log_Pb = [0] * (d + 1) # Target length is d+1 for degree d polynomial
for b, count in counts.items():
if count == 0: continue # Skip if count is zero
# Construct P_b(x) = sum_{k=0..b} x^k / k! up to degree d
Pb = [0] * (d + 1)
for k in range(min(b, d) + 1): # Fill coefficients up to degree min(b, d)
Pb[k] = inv_fact[k]
# Compute log(P_b(x)) mod x^(d+1)
log_Pb = poly_log(Pb, d + 1)
# Add C_b * log(P_b(x)) to the total log sum
for i in range(d + 1):
sum_log_Pb[i] = (sum_log_Pb[i] + count * log_Pb[i]) % MOD
# Compute F(x) = exp(sum_log_Pb) mod x^(d+1)
# F(x) = Product_{i=1..N} P_{B_i}(x) = Product_b P_b(x)^{C_b}
F = poly_exp(sum_log_Pb, d + 1)
# Extract coefficient [x^d] F(x)
coeff_d = F[d] if d < len(F) else 0
# Compute S_d = d! * [x^d] F(x) / d! = d! * coeff_d
Sd = (fact[d] * coeff_d) % MOD
# Calculate phi(K/d)
phi_val = get_phi(K // d)
# Add contribution phi(K/d) * S_d to the total sum
total_S = (total_S + phi_val * Sd) % MOD
# Final answer is (1/K) * total_S mod MOD
K_inv = pow(K, MOD - 2, MOD)
final_ans = (total_S * K_inv) % MOD
print(final_ans)
solve()