結果

問題 No.1067 #いろいろな色 / Red and Blue and more various colors (Middle)
ユーザー lam6er
提出日時 2025-04-09 21:05:55
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 5,577 bytes
コンパイル時間 174 ms
コンパイル使用メモリ 82,284 KB
実行使用メモリ 85,420 KB
最終ジャッジ日時 2025-04-09 21:08:01
合計ジャッジ時間 5,705 ms
ジャッジサーバーID
(参考情報)
judge5 / judge2
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 1
other AC * 10 TLE * 1 -- * 14
権限があれば一括ダウンロードができます

ソースコード

diff #

import bisect

MOD = 998244353

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    ptr = 0
    N, Q = int(data[ptr]), int(data[ptr+1])
    ptr +=2
    A = list(map(int, data[ptr:ptr+N]))
    ptr +=N
    queries = []
    max_A = max(A)
    for _ in range(Q):
        l = int(data[ptr])
        r = int(data[ptr+1])
        p = int(data[ptr+2])
        queries.append((l, r, p))
        ptr +=3
    
    # Preprocessing: sorted A, prefix_prod_notS, suffix_prod_S
    sorted_A = sorted(A)
    prefix_prod_notS = [1]*(N+1)  # prefix_prod_notS[i] = product of first i elements (sorted)
    for i in range(N):
        prefix_prod_notS[i+1] = prefix_prod_notS[i] * sorted_A[i] % MOD
    
    suffix_prod_S = [1]*(N+1)  # suffix_prod_S[i] = product of (a_j -1) from i to N-1
    for i in range(N-1, -1, -1):
        suffix_prod_S[i] = suffix_prod_S[i+1] * (sorted_A[i]-1) % MOD
    
    # Precompute factorial and inverse for combination
    max_n = 6000
    factorial = [1]*(max_n+1)
    for i in range(1, max_n+1):
        factorial[i] = factorial[i-1] * i % MOD
    inv_fact = [1]*(max_n+1)
    inv_fact[max_n] = pow(factorial[max_n], MOD-2, MOD)
    for i in range(max_n-1, -1, -1):
        inv_fact[i] = inv_fact[i+1] * (i+1) % MOD
    
    def comb(n, k):
        if n <0 or k <0 or k >n:
            return 0
        return factorial[n] * inv_fact[k] % MOD * inv_fact[n -k] % MOD
    
    for l, r, p in queries:
        ans = 0
        for c in range(l, r+1):
            left = bisect.bisect_left(sorted_A, c)
            K = N - left
            if K < p:
                ans ^= 0
                continue
            
            # Check if any a_i == c in S (sorted_A[left ... N-1])
            # Find the number of elements ==c in this range
            # Using bisect to find left_c and right_c
            left_c = bisect.bisect_left(sorted_A, c)
            right_c = bisect.bisect_right(sorted_A, c)
            count_c = right_c - left_c
            if left >= right_c:
                count_c =0
            else:
                count_c = max(0, right_c - max(left, left_c))
            
            # Check if c ==1 and a_i ==1
            if c ==1 and count_c >0:
                m = count_c
                if p < m:
                    ans ^=0
                    continue
                new_p = p - m
                T2_count = K - count_c
                if new_p > T2_count:
                    ans ^=0
                    continue
                # Compute sigma_new_p for T2: elements in S and a_i >c (which is 1)
                # T2 is the elements in S (sorted_A[left:]) that are >=c, a_i> c (so a_i >=c+1)
                # but since sorted_A is sorted, a_i> c => a_i >=c+1.
                # T2_count is K - m
                # So in the S, the T2 is elements from (right_c ... N-1)
                T2_start = right_c
                T2_len = N - T2_start
                if T2_len < new_p:
                    ans ^=0
                    continue
                # Extract x_i for T2: sorted_A[T2_start ... N-1], x_i=1/(a_i-1)
                # Compute sigma new_p
                dp = [0]*(new_p+1)
                dp[0] =1
                for i in range(T2_start, N):
                    ai = sorted_A[i]
                    x = pow(ai-1, MOD-2, MOD)
                    for j in range(min(new_p, i - T2_start +1), 0, -1):
                        dp[j] = (dp[j] + dp[j-1] * x) % MOD
                sigma = dp[new_p]
                # Compute comb(T2_len, new_p)
                c_val = comb(T2_len, new_p)
                # prod_S_T2: product of (a_i-1) for T2
                prod_S_T2 = suffix_prod_S[T2_start] if T2_start <= N else 1
                # prod_notS: product of A_i for A_i <c
                prefix_part = prefix_prod_notS[left]  # since left <= T2_start (because T2_start >= right_c >= left)
                # Wait, no: elements <c are up to left-1 (bisect_left)
                # so prefix_prod_notS[left]
                prefix_val = prefix_prod_notS[left] if left >=0 else 1
                res = prefix_val * prod_S_T2 % MOD
                res = res * c_val % MOD
                res = res * sigma % MOD
                ans ^= res
            else:
                # Normal case: all elements in S have a_i >=c and a_i !=c or c !=1, so a_i-1 >=1
                # So compute sigma_p for all S elements (left to N-1)
                # Compute sigma_p as the elementary symmetric sum of x_i=1/(a_i-1)
                K = N - left
                if K < p:
                    ans ^=0
                    continue
                dp = [0]*(p+1)
                dp[0] =1
                valid = True
                for i in range(left, N):
                    ai = sorted_A[i]
                    x = pow(ai-1, MOD-2, MOD) if (ai -1 !=0) else 0
                    if ai -1 ==0:
                        valid = False
                        break
                    for j in range(min(p, i - left +1), 0, -1):
                        dp[j] = (dp[j] + dp[j-1] * x) % MOD
                if not valid:
                    ans ^=0
                    continue
                sigma = dp[p]
                # Compute prod_S = suffix_prod_S[left]
                prod_S = suffix_prod_S[left]
                # Compute prod_notS = prefix_prod_notS[left]
                prod_notS = prefix_prod_notS[left]
                res = prod_notS * prod_S % MOD
                res = res * sigma % MOD
                ans ^= res
        print(ans % MOD)
                
if __name__ == '__main__':
    main()
0