結果

問題 No.1575 Divisor Function Hard
ユーザー lam6er
提出日時 2025-04-09 20:56:33
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,148 bytes
コンパイル時間 437 ms
コンパイル使用メモリ 82,404 KB
実行使用メモリ 132,732 KB
最終ジャッジ日時 2025-04-09 20:58:39
合計ジャッジ時間 9,236 ms
ジャッジサーバーID
(参考情報)
judge5 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 7 TLE * 1 -- * 44
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
from collections import defaultdict

MOD = 998244353

def main():
    input = sys.stdin.read().split()
    ptr = 0
    N, M, K, P, Q = map(int, input[ptr:ptr+5])
    ptr +=5
    A = list(map(int, input[ptr:ptr+N]))
    ptr +=N
    B = list(map(int, input[ptr:ptr+M]))
    ptr +=M
    C = list(map(int, input[ptr:ptr+K]))
    ptr +=K
    U_list = list(map(int, input[ptr:ptr+Q]))
    ptr +=Q

    max_U = max(U_list) if U_list else 0
    max_A = max(A) if A else 0

    # Precompute count_A[i]
    count_A = [0] * (max_A + 2)
    for a in A:
        count_A[min(a, max_A)] += 1
    for i in range(max_A - 1, 0, -1):
        count_A[i] += count_A[i + 1]
    for i in range(max_A + 1):
        count_A[i] %= MOD

    # Precompute S_B[s] and S_C[t]
    def compute_power_sums(arr, max_p):
        mod_arr = [x % MOD for x in arr]
        cnt = defaultdict(int)
        for x in mod_arr:
            cnt[x] += 1
        unique_x = list(cnt.keys())
        power_sums = [0] * (max_p + 1)
        for s in range(max_p + 1):
            total = 0
            for x in unique_x:
                total = (total + pow(x, s, MOD) * cnt[x]) % MOD
            power_sums[s] = total
        return power_sums

    S_B = compute_power_sums(B, P)
    S_C = compute_power_sums(C, P)

    # Precompute combination(P, t)
    fact = [1] * (P + 1)
    for i in range(1, P + 1):
        fact[i] = fact[i-1] * i % MOD
    inv_fact = [1] * (P + 1)
    inv_fact[P] = pow(fact[P], MOD-2, MOD)
    for i in range(P-1, -1, -1):
        inv_fact[i] = inv_fact[i+1] * (i+1) % MOD
    comb = [0] * (P + 1)
    for t in range(P + 1):
        comb[t] = fact[P] * inv_fact[t] % MOD * inv_fact[P - t] % MOD

    # Compute a[s] = comb[s] * S_B[s] * S_C[P-s]
    a = [0] * (P + 1)
    for s in range(P + 1):
        c = comb[s]
        t = P - s
        sc = S_C[t] if 0 <= t <= P else 0
        a[s] = c * S_B[s] % MOD
        a[s] = a[s] * sc % MOD

    # Function to evaluate a polynomial with coefficients a at point x
    def eval_poly(x):
        res = 0
        x_pow = 1
        for s in range(P + 1):
            res = (res + a[s] * x_pow) % MOD
            x_pow = x_pow * x % MOD
        return res

    max_k = max_U
    sum_terms = [0] * (max_k + 2)
    for k in range(1, max_k + 1):
        sum_terms[k] = eval_poly(k)

    prefix_sum = [0] * (max_k + 2)
    for k in range(1, max_k + 1):
        prefix_sum[k] = (prefix_sum[k-1] + sum_terms[k]) % MOD

    # Precompute i^P mod MOD for 1 <= i <= max_A
    pow_iP = [0] * (max_A + 2)
    for i in range(1, max_A + 1):
        pow_iP[i] = pow(i, P, MOD)

    # Process each query
    results = []
    for U in U_list:
        res = 0
        max_i = min(U, max_A)
        for i in range(1, max_i + 1):
            if count_A[i] == 0:
                continue
            m = U // i
            if m == 0:
                continue
            m = min(m, max_k)
            term = count_A[i] * pow_iP[i] % MOD
            term = term * prefix_sum[m] % MOD
            res = (res + term) % MOD
        results.append(res % MOD)

    for result in results:
        print(result)

if __name__ == '__main__':
    main()
0