結果

問題 No.2025 Select $k$-th Submultiset
ユーザー lam6er
提出日時 2025-04-16 00:13:00
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 339 ms / 2,000 ms
コード長 5,101 bytes
コンパイル時間 182 ms
コンパイル使用メモリ 81,452 KB
実行使用メモリ 106,928 KB
最終ジャッジ日時 2025-04-16 00:14:19
合計ジャッジ時間 10,718 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 42
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

def main():
    input = sys.stdin.read().split()
    ptr = 0
    N = int(input[ptr])
    ptr += 1
    L = int(input[ptr])
    ptr += 1
    c = list(map(int, input[ptr:ptr+N]))
    ptr += N
    Q = int(input[ptr])
    ptr += 1
    queries = []
    for _ in range(Q):
        queries.append(int(input[ptr]))
        ptr += 1

    # Precompute DP tables and prefix arrays
    max_L = L
    dp = [ [0]*(max_L +1) for _ in range(N+2) ]  # 1-based to N
    prefix_arrays = [None]*(N+2)  # prefix_arrays[i] is for dp[i+1]

    for i in range(N, 0, -1):
        if i == N:
            for s in range(0, max_L +1):
                if s <= c[i-1]:
                    dp[i][s] = 1
                else:
                    dp[i][s] = 0
        elif i == N-1:
            next_i = i +1
            c_i = c[i-1]
            c_next = c[next_i-1]
            for s in range(0, max_L +1):
                low = max(0, s - c_next)
                high = min(c_i, s)
                if low > high:
                    dp[i][s] = 0
                else:
                    dp[i][s] = high - low +1
        else:
            next_i = i +1
            # Compute prefix sum for dp[next_i]
            prefix = [0]*(max_L +1)
            prefix[0] = dp[next_i][0]
            for s in range(1, max_L +1):
                prefix[s] = prefix[s-1] + dp[next_i][s]
            # Compute dp[i][s]
            c_i = c[i-1]
            for s in range(0, max_L +1):
                a_max = min(c_i, s)
                lower = s - a_max
                if lower <0:
                    lower =0
                if lower ==0:
                    sum_val = prefix[s] if s <= max_L else 0
                else:
                    if s > max_L:
                        sum_val =0
                    else:
                        sum_val = prefix[s]
                        if lower-1 >=0:
                            sum_val -= prefix[lower-1]
                dp[i][s] = sum_val

        # Compute prefix array for i if i < N
        if i < N:
            next_i = i +1
            prefix_arrays[i] = [0]*(max_L +1)
            prefix_arrays[i][0] = dp[next_i][0]
            for s in range(1, max_L +1):
                prefix_arrays[i][s] = prefix_arrays[i][s-1] + dp[next_i][s]

    # Process queries
    total = dp[1][L]
    for k in queries:
        if k > total:
            print(-1)
            continue
        remaining_k = k
        remaining_sum = L
        result = [0]*N
        found = True
        for current_i in range(N):
            i = current_i +1  # 1-based
            s = remaining_sum
            if s <0:
                found = False
                break
            if i == N:
                # last element
                if s > c[current_i] or remaining_k !=1:
                    found = False
                    break
                result[current_i] = s
                remaining_sum -= s
                remaining_k =0
                break
            # Binary search for a_i
            x_max = min(c[current_i], s)
            low =0
            high = x_max
            ans = -1
            prefix = prefix_arrays[i]
            while low <= high:
                mid = (low + high) //2
                lower_t = s - x_max
                upper_t = s - mid
                if upper_t <0:
                    sum_ways =0
                else:
                    if lower_t <0:
                        lower_t =0
                    if upper_t > max_L:
                        sum_ways =0
                    else:
                        sum_ways = prefix[upper_t]
                        if lower_t >0:
                            if lower_t-1 > max_L:
                                sum_ways -=0
                            else:
                                sum_ways -= prefix[lower_t -1]
                if sum_ways >= remaining_k:
                    ans = mid
                    low = mid +1
                else:
                    high = mid -1
            if ans == -1:
                found = False
                break
            # Compute sum_ways_gt
            sum_ways_gt =0
            if ans +1 <= x_max:
                lower_t_gt = s - x_max
                upper_t_gt = s - (ans +1)
                if upper_t_gt <0:
                    sum_ways_gt =0
                else:
                    if lower_t_gt <0:
                        lower_t_gt =0
                    if upper_t_gt > max_L:
                        sum_ways_gt =0
                    else:
                        sum_ways_gt = prefix[upper_t_gt]
                        if lower_t_gt >0:
                            if lower_t_gt -1 > max_L:
                                sum_ways_gt -=0
                            else:
                                sum_ways_gt -= prefix[lower_t_gt -1]
            remaining_k -= sum_ways_gt
            result[current_i] = ans
            remaining_sum -= ans
        if found and remaining_sum ==0 and remaining_k ==0:
            print(' '.join(map(str, result)))
        else:
            print(-1)

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