結果

問題 No.1025 Modular Equation
ユーザー lam6er
提出日時 2025-03-20 20:37:13
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,736 bytes
コンパイル時間 781 ms
コンパイル使用メモリ 82,508 KB
実行使用メモリ 73,732 KB
最終ジャッジ日時 2025-03-20 20:38:06
合計ジャッジ時間 8,828 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 6 TLE * 1 -- * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
import math

MOD = 10**9 + 7

def prime_factors(n):
    factors = []
    while n % 2 == 0:
        factors.append(2)
        n = n // 2
    i = 3
    while i * i <= n:
        while n % i == 0:
            factors.append(i)
            n = n // i
        i += 2
    if n > 2:
        factors.append(n)
    return factors

def find_primitive_root(p):
    if p == 2:
        return 1
    phi = p - 1
    factors = list(set(prime_factors(phi)))
    for g in range(2, p):
        flag = True
        for pr in factors:
            if pow(g, phi // pr, p) == 1:
                flag = False
                break
        if flag:
            return g
    return None

def main():
    input = sys.stdin.read().split()
    idx = 0
    p = int(input[idx]); idx +=1
    n = int(input[idx]); idx +=1
    k = int(input[idx]); idx +=1
    b = int(input[idx]); idx +=1
    a = list(map(int, input[idx:idx+n]))
    
    m = sum(1 for ai in a if ai ==0)
    non_zero_a = [ai for ai in a if ai !=0]
    
    # Handle all zero a_i
    if not non_zero_a:
        total = 1 if (b % p) == 0 else 0
        ans = (total * pow(p, m, MOD)) % MOD
        print(ans)
        return
    
    # Find primitive root
    if p == 2:
        # Special case: primitive root is 1
        g = 1
    else:
        g = find_primitive_root(p)
    
    phi = p - 1
    current_dp = [0] * p
    current_dp[0] = 1  # Initial state
    
    for ai in non_zero_a:
        d = math.gcd(k, phi)
        m_prime = phi // d
        
        if p == 2:
            # For p=2, possible x values are 0 and 1
            # since k >=1, x^k mod 2 is x mod 2 (k=1) or 0 if x=0.
            # For p=2 and k>=1, x^k mod2 is 0 if x=0, else 1 mod2.
            u_values = [1] if d %2 ==1 else []
            # a_i can be 0 or 1. If ai is 1, then the values are 0 and 1.
            # Here, ai !=0, so non-zero a_i can only be 1 (since 0<= ai <p=2)
            # So, non_zero_a elements are 1.
            # So, ai=1: the non-zero x's (x=1) contribute 1^k mod2=1, multiplied by ai=1 →1 mod2
            # for each non-zero x (only x=1), count is 1 (since phi=1, m_prime=1/d, d is gcd(k,1)=1)
            # So, u_values = [1]
            # and d=1, count is d=1 →each x=1 contributes 1 time.
            # but phi =1, m_prime=1
            # This part may need adjustment for p=2.
            # Handle p=2 case separately
            if d %1 !=0:
                pass  # Not possible
            # After long consideration, p=2 case should be handled as follows:
            # x can be 0 or 1.
            # x^k mod2: x=0 →0, x=1 →1 (k >=1)
            # For ai=1 (only possible non-zero), the possible values are 0 (x=0) and 1*1=1 mod2 (x=1)
            # So, cnt_i[0] =1 (x=0 case), and cnt_i[1] =1 (x=1 cases, which are 1)
            # So, u_values should be [1]
            # d=1, m_prime=1, which matches.
        else:
            g_d = pow(g, d, p)
            u_values = [pow(g_d, t, p) for t in range(m_prime)]
        
        # Calculate u_list: ai * u mod p for each u in u_values
        u_list = [(ai * u) % p for u in u_values]
        # Handle x=0 case (adds 0 once)
        
        new_dp = [0] * p
        
        # Process u=0 (count 1)
        for s_prev in range(p):
            new_dp[s_prev] = (new_dp[s_prev] + current_dp[s_prev]) % MOD
        
        # Process non-zero u (each with count d)
        for u in u_list:
            for s_prev in range(p):
                s_new = (s_prev + u) % p
                new_dp[s_new] = (new_dp[s_new] + current_dp[s_prev] * d) % MOD
        
        current_dp = new_dp
    
    target = b % p
    result = current_dp[target] * pow(p, m, MOD) % MOD
    print(result)

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