結果

問題 No.2025 Select $k$-th Submultiset
ユーザー ZHU QIHAO
提出日時 2025-05-15 02:06:49
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 539 ms / 2,000 ms
コード長 4,253 bytes
コンパイル時間 544 ms
コンパイル使用メモリ 82,304 KB
実行使用メモリ 104,900 KB
最終ジャッジ日時 2025-05-15 02:07:09
合計ジャッジ時間 17,245 ms
ジャッジサーバーID
（参考情報） 		judge5 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 42
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

def main():
    data = sys.stdin.read().strip().split()
    if not data:
        return
    it = iter(data)
    N = int(next(it)); L_val = int(next(it))
    c = [int(next(it)) for _ in range(N)]
    Q = int(next(it))
    queries = [int(next(it)) for _ in range(Q)]

    L = L_val
    dp = [[0] * (L + 1) for _ in range(N + 2)]
    prefix = [[0] * (L + 1) for _ in range(N + 2)]

    dp[N + 1][0] = 1
    for j in range(0, L + 1):
        if j == 0:
            prefix[N + 1][j] = dp[N + 1][j]
        else:
            prefix[N + 1][j] = prefix[N + 1][j - 1] + dp[N + 1][j]

    for i in range(N, 0, -1):
        current_c = c[i - 1]
        for j in range(0, L + 1):
            a_max = min(current_c, j)
            low_bound_index = j - a_max - 1

            if low_bound_index < 0:
                dp[i][j] = prefix[i + 1][j] if j <= L else prefix[i + 1][L]
            else:
                if j > L:
                    part1 = prefix[i + 1][L]
                else:
                    part1 = prefix[i + 1][j]
                if low_bound_index > L:
                    part2 = prefix[i + 1][L]
                else:
                    part2 = prefix[i + 1][low_bound_index]
                dp[i][j] = part1 - part2

        for j in range(0, L + 1):
            if j == 0:
                prefix[i][j] = dp[i][j]
            else:
                prefix[i][j] = prefix[i][j - 1] + dp[i][j]

    out_lines = []

    for k_query in queries:
        remaining = L
        current_k = k_query
        freq = [0] * (N + 1)
        valid = True

        for i in range(1, N + 1):
            a_max = min(c[i - 1], remaining)
            total_ways = dp[i][remaining]
            if current_k > total_ways:
                valid = False
                break

            low_bound_index = remaining - a_max - 1

            low = 0
            high = a_max + 1
            candidate_x = 0

            while low <= high:
                mid = (low + high) // 2
                if mid == 0:
                    cum_mid = 0
                else:
                    end_index = remaining - a_max + mid - 1
                    end_index = min(end_index, L)
                    end_index = max(end_index, -1)

                    if end_index < 0:
                        cum_end_index = 0
                    else:
                        cum_end_index = prefix[i + 1][end_index]

                    if low_bound_index < 0:
                        cum_mid = cum_end_index
                    else:
                        cum_low_index = min(low_bound_index, L)
                        cum_low_index = max(cum_low_index, -1)
                        if cum_low_index < 0:
                            cum_low_val = 0
                        else:
                            cum_low_val = prefix[i + 1][cum_low_index]
                        cum_mid = cum_end_index - cum_low_val

                if cum_mid < current_k:
                    candidate_x = mid
                    low = mid + 1
                else:
                    high = mid - 1

            candidate = a_max - candidate_x

            if candidate_x == 0:
                cum_candidate = 0
            else:
                end_idx = remaining - a_max + candidate_x - 1
                end_idx = min(end_idx, L)
                end_idx = max(end_idx, -1)

                if end_idx < 0:
                    val1 = 0
                else:
                    val1 = prefix[i + 1][end_idx]

                if low_bound_index < 0:
                    val2 = 0
                else:
                    low_idx = min(low_bound_index, L)
                    low_idx = max(low_idx, -1)
                    if low_idx < 0:
                        val2 = 0
                    else:
                        val2 = prefix[i + 1][low_idx]

                cum_candidate = val1 - val2

            current_k -= cum_candidate
            freq[i] = candidate
            remaining -= candidate

        if not valid or remaining != 0:
            out_lines.append("-1")
        else:
            out_str = " ".join(str(freq[i]) for i in range(1, N + 1))
            out_lines.append(out_str)

    for line in out_lines:
        print(line)

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