結果

問題 No.2033 Chromatic Duel
ユーザー qwewe
提出日時 2025-05-14 13:23:50
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 94 ms / 2,000 ms
コード長 3,360 bytes
コンパイル時間 185 ms
コンパイル使用メモリ 82,048 KB
実行使用メモリ 80,640 KB
最終ジャッジ日時 2025-05-14 13:25:21
合計ジャッジ時間 3,867 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 4
other AC * 37
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

def solve():
    N, B, W_target = map(int, sys.stdin.readline().split())

    MOD = 998244353

    # Max n for nCr is W_target + num_vars - 1.
    # Max W_target is N. Max num_vars is k1+n2 <= 2 + (B-1) = B+1.
    # So, N + (B+1) - 1 = N+B.
    # Max N+B is 2*10^5 + 2*10^5 = 4*10^5.
    MAX_COMB_N = N + B + 5 
    fact = [1] * (MAX_COMB_N + 1)
    inv_fact = [1] * (MAX_COMB_N + 1)

    for i in range(1, MAX_COMB_N + 1):
        fact[i] = (fact[i-1] * i) % MOD

    if MAX_COMB_N >=0: # Handle MAX_COMB_N can be -1 if N+B+5 is 0, (e.g. N=B=0)
                       # However, N>=3, B>=1, so MAX_COMB_N is always positive.
        inv_fact[MAX_COMB_N] = pow(fact[MAX_COMB_N], MOD - 2, MOD)
    
    for i in range(MAX_COMB_N - 1, -1, -1): 
        inv_fact[i] = (inv_fact[i+1] * (i+1)) % MOD
    
    # memo_nCr = {} # Memoization for nCr can be skipped if direct computation is fast enough
    def nCr_mod(n, r):
        if r < 0 or r > n:
            return 0
        
        num = fact[n]
        den = (inv_fact[r] * inv_fact[n-r]) % MOD
        res = num * den % MOD
        return res

    ans = 0

    target_total_lost_cells = (N - B) - W_target
    # This is guaranteed >= 0 because B + W_target <= N  => W_target <= N - B

    # k1: number of end segments (y_0, y_B) that are non-empty (y_j >= 1, lost_cells = 1)
    # There are 2 end segments. k1 can be 0, 1, or 2.
    for k1 in range(3): 
        coef_k1 = nCr_mod(2, k1)

        num_middle_segments = B - 1
        
        # n2: number of middle segments (y_1 to y_{B-1}) that are y_i >= 2 (lost_cells = 2)
        # Max n2 is num_middle_segments. Loop goes up to num_middle_segments inclusive.
        # If num_middle_segments is -1 (B=0, not possible), range(0) is empty.
        # If num_middle_segments is 0 (B=1), range(1) means n2=0.
        for n2 in range(num_middle_segments + 1 if num_middle_segments >=0 else 0 ): 
            coef_n2 = nCr_mod(num_middle_segments, n2)
                
            lost_by_k1_n2 = k1 * 1 + n2 * 2
            
            # n1: number of middle segments that are y_i = 1 (lost_cells = 1)
            n1 = target_total_lost_cells - lost_by_k1_n2
            
            # These n1 segments must be chosen from remaining_middle_slots
            remaining_middle_slots_for_n1_or_n0 = num_middle_segments - n2
            
            if n1 < 0 or n1 > remaining_middle_slots_for_n1_or_n0:
                continue # This choice of k1, n2 is not possible
                
            coef_n1 = nCr_mod(remaining_middle_slots_for_n1_or_n0, n1)

            num_vars_for_W_dist = k1 + n2
            
            ways_dist_W = 0
            if num_vars_for_W_dist == 0:
                if W_target == 0:
                    ways_dist_W = 1
            else: # num_vars_for_W_dist > 0
                # W_target must be non-negative.
                # Distribute W_target items into num_vars_for_W_dist bins.
                # Stars and bars: C(W_target + num_vars - 1, num_vars - 1)
                ways_dist_W = nCr_mod(W_target + num_vars_for_W_dist - 1, num_vars_for_W_dist - 1)

            term = coef_k1
            term = (term * coef_n2) % MOD
            term = (term * coef_n1) % MOD
            term = (term * ways_dist_W) % MOD
            
            ans = (ans + term) % MOD

    sys.stdout.write(str(ans) + "\n")

solve()
0