Patisserie ABC 2 [京セラプログラミングコンテスト2021（AtCoder Beginner Contest 200） E]

https://atcoder.jp/contests/abc200/tasks/abc200_e

解説

https://atcoder.jp/contests/abc200/submissions/22441425

今回の問題は辞書順でK番目のものを答えよという問題と同様の解き方で解いていく。
辞書順でK番目のものを答えよという問題の場合は

[A.....]
[B.....]
[C.....]

みたいなグループを考えて、そのグループに属する数の個数を計算して、先頭からK番目が含まれるものを探していくようなやり方をしていく。
これを計算するにあたって、

dp[cnt][tot] := 1～Nのいずれかであるcnt個の数を使って総和がtotとなる組合せ

が必要になる。

dp[cnt][tot]

これを自分はDPでimos法を使った高速化で計算した。
普通に計算すると以下のようになる。

rep(cnt, 0, 3) rep(tot, 1, N + 1) {  
    rep(delta, 1, N + 1) dp[cnt + 1][tot + delta] += dp[cnt][tot];  
}  

これが愚直に書いた場合の更新式になるが、dp[cnt][tot]の値がdp[cnt+1][tot+1]～dp[cnt+1][tot+N]に一様に足されるので、
imos法を用いれば高速化が可能。
これによって高速化する。

最初から選択していく

前から見ていくと

[総和が3のグループ]
[総和が4のグループ]
…

みたいになっているので、総和を3から順番に見ていって、個数を見ながらK番目がどのグループに入るかを確認していく。
総和がtotのグループに入っている個数はdp[3][tot]個となる。

総和が特定できたら、次は綺麗さがどれかを特定していく。
とある総和のグループの中身を細かく見ていると

[綺麗さが1のグループ]
[綺麗さが2のグループ]
…

みたいになっている。
これも同様に小さい方から見ていく。
綺麗さが1のグループに入っている個数はdp[2][tot-i]個となる。

最後においしさを特定する。
これも同様に

[美味しさ1のグループ]
[美味しさ2のグループ]
…

みたいになってるので、dp[1][tot-i-j]個であることを利用して特定する。
これで答える。

int N; ll K;
ll dp[4][3010101];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K;

    dp[0][0] = 1;
    rep(cnt, 0, 3) {
        rep(tot, 0, 2010101) {
            dp[cnt + 1][tot + 1] += dp[cnt][tot];
            dp[cnt + 1][tot + N + 1] -= dp[cnt][tot];
        }
        rep(tot, 1, 3010101) {
            dp[cnt + 1][tot] += dp[cnt + 1][tot - 1];
            chmin(dp[cnt + 1][tot], infl);
        }
    }

    ll cu = 1;
    int tot = -1;
    rep(t, 3, 3010101) {
        if (cu <= K && K < cu + dp[3][t]) {
            tot = t;
            break;
        }
        cu += dp[3][t];
    }

    rep(t, 0, tot) K -= dp[3][t];

    cu = 1;
    int ans_i = -1;
    rep(i, 1, N + 1) {
        if (cu <= K && K < cu + dp[2][tot - i]) {
            ans_i = i;
            break;
        }
        cu += dp[2][tot - i];
    }

    rep(t, 1, ans_i) K -= dp[2][tot - t];

    cu = 1;
    int ans_j = -1;
    rep(j, 1, N + 1) {
        if (cu <= K && K < cu + dp[1][tot - ans_i - j]) {
            ans_j = j;
            break;
        }
        cu += dp[1][tot - ans_i - j];
    }

    int ans_k = tot - ans_i - ans_j;
    printf("%d %d %d\n", ans_i, ans_j, ans_k);
}