結果

問題 No.2025 Select $k$-th Submultiset
ユーザー hahho
提出日時 2025-02-10 21:47:30
言語 C++17(gcc12)
(gcc 12.3.0 + boost 1.87.0)
結果
AC  
実行時間 97 ms / 2,000 ms
コード長 6,431 bytes
コンパイル時間 3,408 ms
コンパイル使用メモリ 114,096 KB
実行使用メモリ 12,672 KB
最終ジャッジ日時 2025-02-10 21:47:41
合計ジャッジ時間 10,128 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 42
権限があれば一括ダウンロードができます

ソースコード

diff # 

// ChatGPT o3-mini-highのコード

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

const long long INF = 1000000000000000001LL;

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N;
    long long L;
    cin >> N >> L;
    // 1-indexed: c[1]～c[N]
    vector<long long> c(N+1);
    for (int i = 1; i <= N; i++){
        cin >> c[i];
    }
    
    // 特にN==1の場合の処理
    if(N == 1){
        int Q;
        cin >> Q;
        for (int qi = 0; qi < Q; qi++){
            long long k;
            cin >> k;
            // Sは数字1のみからなり、c[1]個ある。
            // 長さLの部分列が存在するのは L <= c[1] で、可能な部分列はただ一通り（L個の1）。
            if(L <= c[1] && k == 1){
                cout << L << "\n";
            } else {
                cout << -1 << "\n";
            }
        }
        return 0;
    }
    
    // Smax[i] = c[i] + c[i+1] + ... + c[N] を計算（i=1～N, かつ Smax[N+1]=0）
    vector<int> Smax(N+2, 0);
    Smax[N+1] = 0;
    for (int i = N; i >= 1; i--){
        Smax[i] = c[i] + Smax[i+1];
    }
    
    // dp[i][u] = 「数字 i～N から、選ぶ個数の和が u となる採り方の個数」
    // ※ dp[i]の添字は 0～Smax[i] まで。ここでは i=2～N の分を前計算する（digit1はクエリ処理中に貪欲に決定）
    vector< vector<long long> > dp(N+2);
    vector< vector<long long> > cum(N+2); // dpの累積和
    
    // i = N の場合: Sに数字Nが c[N] 個あるので、部分列として選べる個数は0～c[N]のみで、
    // 各 u (0<=u<=c[N]) で採り方は1通り、その他は0通り
    dp[N].resize(Smax[N] + 1, 0);
    cum[N].resize(Smax[N] + 1, 0);
    for (int u = 0; u <= c[N]; u++){
        dp[N][u] = 1;
    }
    // 累積和 cum[N]
    cum[N][0] = dp[N][0];
    for (int u = 1; u <= Smax[N]; u++){
        long long tmp = cum[N][u-1] + dp[N][u];
        if(tmp > INF) tmp = INF;
        cum[N][u] = tmp;
    }
    
    // i = N-1 から 2 まで：dp[i][u] = sum_{x=max(0, u-Smax[i+1])}^{min(c[i], u)} dp[i+1][u-x]
    for (int i = N-1; i >= 2; i--){
        dp[i].resize(Smax[i] + 1, 0);
        cum[i].resize(Smax[i] + 1, 0);
        for (int u = 0; u <= Smax[i]; u++){
            // digit i から選ぶ個数 x の候補は、
            // x ∈ [max(0, u - Smax[i+1]), min(c[i], u)]
            int x_low  = max(0, u - Smax[i+1]);
            int x_high = min((int)c[i], u);
            if(x_low > x_high){
                dp[i][u] = 0;
                continue;
            }
            // x を変数変換: w = u - x とすると、wは u - x_high から u - x_low まで
            int Lw = u - x_high;
            int Rw = u - x_low;
            // dp[i+1] は添字 0～Smax[i+1] なので、Rwは Smax[i+1] を超えないはずですが安全のため
            if(Rw > Smax[i+1]) Rw = Smax[i+1];
            long long ways = 0;
            if(Lw == 0) ways = cum[i+1][Rw];
            else ways = cum[i+1][Rw] - cum[i+1][Lw - 1];
            if(ways > INF) ways = INF;
            dp[i][u] = ways;
        }
        // dp[i] の累積和 cum[i]
        cum[i][0] = dp[i][0];
        for (int u = 1; u <= Smax[i]; u++){
            long long tmp = cum[i][u-1] + dp[i][u];
            if(tmp > INF) tmp = INF;
            cum[i][u] = tmp;
        }
    }
    
    // クエリ処理：各クエリで、Sから長さLの部分列（同じ列は重複して数えない）の中で、
    // 辞書順（「d1, d2, ..., dN」について、d1が大きいほど小さい、d1が同じならd2が大きいほど小さい、…）で
    // k番目のものを求め、その各数字の出現個数 d[1]...d[N] を出力する。
    // ※ digit1 については、残りの採用個数 rem と、digit2～Nの通り数（dp[2]）から貪欲に決定する
    int Q;
    cin >> Q;
    for (int qi = 0; qi < Q; qi++){
        long long k;
        cin >> k;
        vector<long long> d(N+1, 0);  // 1-indexed: d[1]...d[N]
        long long rem = L; // 残り、digit2～N で埋めるべき個数も含めた全体の採用個数
        bool ok = true;
        // digit1～digit(N-1)について決定
        for (int i = 1; i <= N-1; i++){
            // digit i で選べる個数 x は、
            // x ∈ [lb, ub]  で lb = max(0, rem - Smax[i+1]), ub = min(c[i], rem)
            int lb = max(0, (int)(rem - Smax[i+1]));
            int ub = min((int)c[i], (int)rem);
            if(lb > ub){ ok = false; break; }
            // x と置換して u = rem - x とおくと、
            // x が大きい（＝ digit i を多く使う）ほど u は小さくなるので、候補 u は [rem-ub, rem-lb] となる。
            int L0 = rem - ub;
            int R0 = rem - lb;
            // dp[i+1] の累積和 cum[i+1]（添字0～Smax[i+1]）を用いて、
            // 区間 [L0, U] の和が k 以上となる最小の U を二分探索する
            int lo_bin = L0, hi_bin = R0, U_found = R0 + 1;
            while(lo_bin <= hi_bin){
                int mid = (lo_bin + hi_bin) / 2;
                long long sumInterval = (L0 == 0 ? cum[i+1][mid] : cum[i+1][mid] - cum[i+1][L0 - 1]);
                if(sumInterval >= k){
                    U_found = mid;
                    hi_bin = mid - 1;
                } else {
                    lo_bin = mid + 1;
                }
            }
            if(U_found > R0){ ok = false; break; }
            // 対応する x = rem - U_found
            int chosen_x = rem - U_found;
            d[i] = chosen_x;
            // さらに、[L0, U_found-1] の累積和分だけ k を減らす
            long long sumIntervalBefore = 0;
            if(U_found - 1 >= L0){
                sumIntervalBefore = (L0 == 0 ? cum[i+1][U_found - 1] : cum[i+1][U_found - 1] - cum[i+1][L0 - 1]);
            }
            k -= sumIntervalBefore;
            rem -= chosen_x;
        }
        // 最後、digit N では d[N] = rem でなければならない（ただし rem <= c[N] である必要がある）
        if(rem < 0 || rem > c[N]) ok = false;
        else d[N] = rem;
        
        if(!ok || k != 1) {
            cout << -1 << "\n";
        } else {
            for (int i = 1; i <= N; i++){
                cout << d[i] << (i == N ? "\n" : " ");
            }
        }
    }
    
    return 0;
}
0