ABC200-E Patisserie ABC 2

atcoder.jp
ケーキをパラメータの合計値でグループ分けする、という考察までは進んだが、そこで詰まった。

【問題】
1<=a,b,c<=Nである(a,b,c)の数列が合計値->各パラメータの昇順に並んでいる。K番目の(a,b,c)の値を出力せよ。
【典型ポイント】
辞書順Ｋ番目を求める問題は、各パラメータごとにグループ分けして考える(よくDPが使える)
【知らなかったこと】
一様の値を一定区間に連続して配る時は、DPだろうがimos法が使える
【解法1】
まず、Ｋ番目のa+b+cの合計値を知りたい。
これは、DP + imos法 で求められる。
dp[i][j]:先頭からi個のパラメータを決めたときの、パラメータの合計値がjである場合の数　とすると
遷移式：dp[i+1][j+(1~N)] += dp[i][j]　
愚直に遷移するとO(N^2)かかってしまうので、一様の値を一定範囲に配ることに着目する。
imosで高速遷移を行うとdpテーブルの計算量がO(N)になって間に合う

K番目の合計値は, i=3から昇順にdp[3][i]をKから引いていき、次に引いた場合Kが0以下になる最初のiとなる。
これでKの合計値が求まった。

あとはa,b,cの値を順次確定させていく。
合計値を求めるのと同じようにa,bを昇順にチェックし、どのグループに属しているかを調べる。

【解法2】
立方体を書いてK番目の合計値を求める方法もある。（自分には使いこなせなさそう）
Nを制限しないとき、a+b+cの場合の数の増え方には規則性がある(等差数列っぽい)のだが、図にあらわすとわかりやすい。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pll = pair<long long, long long>;
constexpr char ln =  '\n';
constexpr long long MOD = 1000000007;
constexpr long long INF = 1000000000 + 100;
constexpr long long LINF = 1000000000000000000 + 100;
#define all(v) v.begin( ), v.end()
#define rep(i,n) for(int i=0;i<(n);i++)
#define rept(i, j, n) for(int i=(j); i<(n); i++)
#define rrep(i, n) for(int i=(n); i>=0; i--)
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

//max(sum(a,b,c))
const int MAX_SUM = 1e6*3+100;

ll dp[4][MAX_SUM];
int main(){
    ll N, K; cin >> N >> K;

    //dp[i][j]:先頭からi番目までの数字を決定したときの、合計値がjである場合の数
    dp[0][0] = 1;
    rep(i, 3){
        rep(j, MAX_SUM){
            if(dp[i][j]==0)continue;
            //imos 前処理
            dp[i+1][j+1] += dp[i][j];
            dp[i+1][j+N+1] -= dp[i][j];
        }
        //imos 同一テーブル上でメモから値を導く
        rept(j,1, MAX_SUM){
            dp[i+1][j] = dp[i+1][j] + dp[i+1][j-1];
        }
    }

    int sum = -1;
    for(int i=3; i<=MAX_SUM; i++){
        if(K <= dp[3][i]){
            sum = i;
            break;
        }
        K -= dp[3][i];
    }
    int ra = -1;
    for(int i=1; i<=N; i++){
        if(K <= dp[2][sum-i]){
            ra = i;
            break;
        }
        K -= dp[2][sum-i];
    }
    int rb = -1;
    for(int i=1; i<=N; i++){
        if(K <= dp[1][sum-ra-i]){
            rb = i;
            break;
        }
        K -= dp[1][sum-ra-i];
    }
    cout << ra << " " << rb << " " << sum-ra-rb << ln;
}