結果

問題 No.2025 Select $k$-th Submultiset
ユーザー lam6er
提出日時 2025-04-16 00:08:17
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 374 ms / 2,000 ms
コード長 3,309 bytes
コンパイル時間 558 ms
コンパイル使用メモリ 81,140 KB
実行使用メモリ 104,048 KB
最終ジャッジ日時 2025-04-16 00:09:19
合計ジャッジ時間 12,368 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 42
権限があれば一括ダウンロードができます

ソースコード

diff # 

def main():
    import sys
    input = sys.stdin.read().split()
    ptr = 0
    N, L = int(input[ptr]), int(input[ptr+1])
    ptr +=2
    c = list(map(int, input[ptr:ptr+N]))
    ptr +=N
    Q = int(input[ptr])
    ptr +=1
    queries = [int(input[ptr+i]) for i in range(Q)]
    
    # Initialize DP and prefix sums
    max_L = L
    dp = [ [0]*(max_L+1) for _ in range(N+2) ]  # dp[i][rem_L]
    prefix_sums = [ [0]*(max_L+2) for _ in range(N+2) ]  # prefix_sums[i][s] = sum(dp[i][0..s-1])
    
    # Base case: i = N+1
    dp[N+1][0] = 1
    for s in range(0, max_L+1):
        prefix_sums[N+1][s+1] = prefix_sums[N+1][s] + dp[N+1][s]
    
    for i in range(N, 0, -1):
        # Compute prefix sums for i+1
        for s in range(0, max_L+1):
            prefix_sums[i+1][s+1] = prefix_sums[i+1][s] + dp[i+1][s]
        # Compute dp[i][rem_L]
        ci = c[i-1]
        for rem_L in range(0, max_L+1):
            max_a = min(ci, rem_L)
            lower = rem_L - max_a
            if lower < 0:
                lower = 0
            upper = rem_L  # a_i=0: rem_L -0 = rem_L
            # sum from s=lower to upper of dp[i+1][s]
            sum_ways = prefix_sums[i+1][upper+1] - prefix_sums[i+1][lower]
            dp[i][rem_L] = sum_ways
        # Compute prefix sums for i
        for s in range(0, max_L+1):
            prefix_sums[i][s+1] = prefix_sums[i][s] + dp[i][s]
    
    # Process queries
    results = []
    for k in queries:
        if k > dp[1][L]:
            results.append("-1")
            continue
        res = [0]*N
        rem = L
        current_k = k
        valid = True
        for i in range(1, N+1):
            ci = c[i-1]
            max_a = min(ci, rem)
            if max_a <0:
                valid = False
                break
            # Binary search for the best a_i
            low = 0
            high = max_a
            best_a = -1
            while low <= high:
                mid = (low + high) // 2
                a_candidate = mid
                # Compute sum_ways for a_i >= a_candidate
                lower_sum = rem - max_a  # rem - max_a = rem - min(ci, rem)
                upper_sum = rem - a_candidate
                if lower_sum <0:
                    lower_sum =0
                if upper_sum <0:
                    sum_ways =0
                else:
                    sum_ways = prefix_sums[i+1][upper_sum +1] - prefix_sums[i+1][lower_sum]
                if sum_ways >= current_k:
                    best_a = a_candidate
                    low = mid +1
                else:
                    high = mid -1
            if best_a == -1:
                valid = False
                break
            # Compute sum_ways for a_i > best_a
            a_gt = best_a +1
            sum_gt =0
            if a_gt <= max_a:
                lower_gt = rem - max_a
                upper_gt = rem - a_gt
                if upper_gt >= lower_gt and upper_gt >=0:
                    sum_gt = prefix_sums[i+1][upper_gt +1] - prefix_sums[i+1][lower_gt]
            current_k -= sum_gt
            res[i-1] = best_a
            rem -= best_a
        if valid and rem ==0:
            results.append(' '.join(map(str, res)))
        else:
            results.append("-1")
    
    print('\n'.join(results))

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