Atcoder ドワンゴからの挑戦状第四回 予選 C問題について一部のテストケースでWAになる

Question

質問事項

Atcoderドワンゴからの挑戦状第四回予選C問題を下をコードで解いたのですが、一部のテストケースでWAになってしまいました。方針としてはこの問題の解説に従ったのですが、私の実装面で不備があるのか、または分割数やdpテーブルの定義に不備があるのかわかりません。どなたか教えていただけますでしょうか。

該当のソースコード

//#define _GLIBCXX _DEBUG
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int INF = INT_MAX;
const int MOD = 1e9 + 7;
vector<vector<ll>> P(1001, vector<ll>(1001,0));
ll solve(vector<int> &kill, vector<int> &death){
    deque<pair<int, int>> cnt;    // cnt...キル数で分けたグループの集合
    cnt.push_back({kill[0], 1});  //cnt.first キル数　cnt.second同じキル数の人数
    for(ll i = 1; i < kill.size(); ++i){
        if(cnt.back().first == kill[i]) cnt.back().second++;
        else cnt.push_back({kill[i], 1});
    }
    int toll = 0;
    for(auto i : death){
        toll += i;
    }
    int n = cnt.size();
    vector<vector<ll>> dp(n + 1, vector<ll>(toll + 1, 0)); // dp[group][d]:= group番目までで、デス数がdである組み合わせ数
    dp[0][0] = 1;
    for(int group = 0; group < n; ++group){
        for(int d = 0; d <= toll; ++d){
            for(int j = 0; j < toll - d + 1; ++j){
                dp[group + 1][d + j] +=  (dp[group][d] * P[cnt[group].second][j]);
                dp[group + 1][d + j] %= MOD;
            }
        }
    }
    return dp[n][toll];
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int N, M; cin >> N >> M;
    vector<int> A(N), B(M);
    for(int i = 0; i < N; ++i) cin >> A[i];
    for(int i = 0; i < M; ++i) cin >> B[i];

    P[0][0] = 1;
    for(int i = 1; i <= 1000; ++i ){
        for(int j = 0; j <= 1000; ++j){
            P[i][j] = P[i - 1][j];
            if(j >= i) P[i][j] += P[i][j - i];  //P[i][j]:= jをi個以下に分割する方法の総数
        }
    }
    ll a = solve(A, B);
    ll b = solve(B, A);
    ll ans = a * b;
    ans %= MOD;
    cout << ans << endl;
}

Answer 1

質問者さんのコードにおいて、分割数の求め方がそもそも正しいかが確認出来ておりませんが、以下の部分においてMODを取り忘れておりませんか？（ここで求めている分割数PはDPテーブルを埋める過程で掛けて使用しているので、あまりに大きいとMODを取る前にオーバーフローしてしまいます）

    for(int i = 1; i <= 1000; ++i ){
        for(int j = 0; j <= 1000; ++j){
            P[i][j] = P[i - 1][j];
      //累積的に計算すると段々と大きな数になりそうなのでオーバーフローを避けるために逐一MODを取る
            if(j >= i) P[i][j] += P[i][j - i];  //P[i][j]:= jをi個以下に分割する方法の総数
        }
    }

もしこのような修正を行ってもWAが取れない場合は、分割数の求め方がそもそも間違っていることが第一の原因として考えられます。質問者さんの参照した解説記事を基に、分割数を求めるアルゴリズムの解説を以下に行いたいと思います。

分割数の分かりやすい言い換え例として、りんごn個をr人に少なくとも1つずつ分ける組み合わせ数(n,rは自然数）を考えましょう。この時、考えられる組み合わせは以下の2通りのいずれかに必ず分類でき、互いに排反（＝両方に当てはまることがない）です。

• 全員が2つ以上のりんごを貰っている（＝1つしか貰えていない人が存在しない）
• 1つしかりんごを貰えていない人が、少なくとも1人存在する

ここで、組み合わせ数を求めようとしていることから、貰う順序を区別していないことに注意して、ここからどのような漸化式を立てられるかを考えてみましょう。
前者に関しては、「n-r個のりんごをr人に配るような全ての組み合わせに、各人+1個ずつりんごを配る」（これによって全員2つ以上のりんごを貰ってることが保証される！）とすれば構成することが出来ます。
後者に関しては、「n-1個のりんごをr-1人に配るような全ての組み合わせに、r人目にりんごを1個配る」（これによって1つしかりんごを貰えていない人の存在が保証される！）とすれば構成することが出来ます。
ここで、新たに構成した組み合わせに関しては、1個のりんごを貰えている人がいるかどうかという点に注目すれば互いに排反であることがわかり、逆に今回考えたいりんごn個をr人に少なくとも1つずつ分ける組み合わせのいかなる場合も、このようにすれば構成出来ることが分かります。
すなわち、この一手前の場合の組み合わせ数を足し合わせたものが、今回考えたい分割数になります。（これは質問者さんの参照した解説記事中の

Q(n,1) = 1 (n≧1), Q(n,n) = 1 (n≧1), Q(n,r) = Q(n-1,r-1) + Q(n-r,r) (n≧r≧2)

に対応します）初期状態Q(n,1) = 1や終了状態Q(n,n) = 1 (n≧1)に注意して、実装例は以下の通りです。

using vll = vector<ll>;
using vvll = vector<vll>;
const int mod1 = 1000000007;
#define REP(i, a, b) for (ll i = a; i < ll(b); i++)

//PN[n][r] := 自然数nをr個の自然数に分割する組み合わせの数（「分割数」）
vvll PN(1010, vll(1010));
void calcPN(int n){
    REP(i,1,n+1) PN[i][1] = 1;

    REP(i,2,n+1){
        REP(j,2,n+1){
            PN[i][j] = PN[i-1][j-1]; //1が少なくとも1つ含まれる場合の数（i個目に1を分けるとみなす）
            if(i - j >= 1){
                PN[i][j] += PN[i-j][j]; //全て2以上の場合の数（PN[i-j][j]の状態の組み合わせからj個全てを+1するとみなす）
                PN[i][j] %= mod1;
            }
        }
    }

    PN[n][n] = 1;
}

ここまでで、自然数nをr個の「自然数」（＝0を含まない）に分割する時の組み合わせの数を求める事ができました。しかし、ここで一つ問題になる点がございます。本問に関しては、自然数nをr個の「非負整数」（＝0を含む）に分割する時の組み合わせの数を知りたいのです。すなわち、「りんごn個をr人に分ける（1つもりんごを貰わない人が存在してもよい）組み合わせ数」が知りたいのです。

この問題に対処するために、質問者さんの参照した解説記事では、自然数nをr個の非負整数(=0も含む）に分割する組み合わせの数を、自然数nをr個以下の自然数に分割する組み合わせの数、とみなすことで、累積的に計算しております。この原理としては、再び先程の例を利用すると、以下のように説明出来ます。
（組み合わせ数を考えるので貰う順序は区別しないことに注意して、）まずは、りんごn個をk人に少なくとも1つずつ分けます。（ここでkはr以下の自然数）。次に、残りのr-k人はりんごを0個貰うようにします。このようにすることで、「りんごn個をr人に分ける（1つもりんごを貰わない人がr-k人存在する）ような組み合わせ」が構成出来ました。後はこのkに関して、1~rの範囲で足し合わせることで、「りんごn個をr人に分ける（1つもりんごを貰わない人が存在してもよい）組み合わせ数」が求まります。

よって、自然数nをr個以下の自然数に分割する組み合わせの数を計算できれば本問に対応出来ることがわかったので、後はこれを実装するだけです。ここでの実装は、
「（自然数nをr個以下の自然数に分割する組み合わせの数）　＝　（自然数nをr-1個以下の自然数に分割する組み合わせの数）+ （自然数nをr個ちょうどの自然数に分割する組み合わせの数）」
という累積和の考え方を用います。実装例は以下の通りです。（質問者さんはこの累積和を元の配列を使い回すことで実装しようとしていたのかもしれませんが、バグの原因になるので、累積和を記録する新たな配列を用意するのがよいでしょう）

//APN[n][r] := 非負整数nをr個の非負整数に分割する組み合わせの数(=非負整数nをr個以下の自然数に分割する組み合わせの数、とみなして累積的に計算出来る)
vvll APN(1010, vll(1010, 0));
void calcAPN(int n){
    rep(i,n+1) APN[0][i] = 1; //例外として、0個をi人に分ける方法を1通りとしている

    REP(i,1,n+1){
        REP(j,1,n+1){
            APN[i][j] = APN[i][j-1] + PN[i][j]; //累積和の計算
            APN[i][j] %= mod1;
        }
    }
}

ここまでで、分割数を求めることが出来たので、後はDPの遷移に適切に組み込むだけです。何か疑問点等ございましたら遠慮なく聞いてください。

Comments

[rdld036]

詳説いただきありがとうございます。
解説の意味がよく理解できました。解説ではmarcさんのおっしゃる通り、りんごn個をr人に少なくとも1つずつ分ける組み合わせ数から考えたと思いますが、私は最初からりんごn個をr人に0個以上分ける組み合わせ数を考えました。主にこちらのサイト
(https://qiita.com/drken/items/f2ea4b58b0d21621bd51)の第6章
を参考にしたのですが、りんごを0個配られる人がいる場合と全員が1個以上配られる場合に分けられます。前者の場合はその0個配られた人を除いたr - 1人にn個りんごを配る場合の数と一致します。後者の場合は全員からりんごを一つずつ取り上げたn - r個からr人に0個以上ずつりんごを配る場合の数と一致します。つまり
P(n,r):=n個のりんごをr人に0個以上配る場合の数
とすると、
P(n,r) = P(n, r - 1) + P(n - r, r)
という漸化式を表せます。つまり上のコード上では
P[r][n] = P[r - 1][n] + P[r][n - r]
です。
そしてmarcさんのおっしゃるとおり、分割数テーブルの作成の際にMODをつけたところ、全てACになりました。

for(int i = 1; i <= 1000; ++i ){
for(int j = 0; j <= 1000; ++j){
P[i][j] = P[i - 1][j];
if(j >= i) P[i][j] += P[i][j - i];
P[i][j] %= MOD;
}
}

[rdld036]

その上でもう一つ質問したいのですが、

for(int group = 0; group < n; ++group){
for(int d = 0; d <= toll; ++d){
for(int j = 0; j < toll - d + 1; ++j){
dp[group + 1][d + j] += (dp[group][d] * P[cnt[group].second][j]);
dp[group + 1][d + j] %= MOD;

この部分で、dpの各要素は足し上げる前はMODをとっているので、dp[group + 1][d + j]は1e9+7以下のはずです。そこにdp[group][d] * P[cnt[group].second][j]を足すのですが、dp[group][d]は1e9 + 7以下、P[cnt[group].second][j]の最大値は実際に

for(int i = 1; i <= 1000; ++i ){
for(int j = 0; j <= 1000; ++j){
P[i][j] = P[i - 1][j];
if(j >= i) P[i][j] += P[i][j - i];
}
}

ll tmp = 0;
int s = 0, t = 0;
for(int i = 0; i <= 1000; ++i){
for(int j = 0; j <= 1000; ++j){
if(tmp < P[i][j]){
s = i; t = j;
tmp = P[i][j];
}
}
}
cout << s << ' ' << t << ' '<< tmp << endl;
cout << tmp * MOD << endl;
cout << LLONG_MAX << endl;

を実行して試したところ、P[170][371]で999999565となりました。
つまりPテーブルの作成の際にMODを取らなかった場合でも、最大でdp[group + 1][d + j]は1e9 + 7 + (1e9 + 7) * 999999565で10^17までしかならないはずです。私の環境でもatcoder の環境でもlong long の最大値は約9 * 10^18でした。なぜそれでもオーバーフローしてしまうのでしょうか。

[marc2825]

これに関しては、質問者さんがP[i][j]を計算する過程で既にオーバーフローしてしまっているため、後半部分でどのi,jで最大になるかを全探索しようとしても正確な値が求まらない（＝その値をもとに最大値を考察しても意味がない）ことが恐らく原因です。実際、「P[i][j]:= jをi個以下に分割する方法の総数」と定義しているのにもかかわらず、(170,371)という中途半端な値で最大値を取っているのは直感にも反しますし、i個「以下」に分割する総数という定義から、iに関して単調増加になることは明らかです。オーバーフローを起こしてしまっていることは、適当に(i,j)に大きな値を代入した時のPを出力して（負の数）等明らかにおかしくなってる部分を発見したり、手計算でも計算できる程シンプルな場合の値と上記のコードでの計算結果を比較したり、オーバーフローを起こさない言語で同じ趣旨のコードを書いて計算結果を比較したりすれば、気づけるはずです。因みに、自分の環境でも質問者さんのコード（MODなし版）を実行してみましたが、そもそもP[170][371]で999999565とならなかったので、やはりオーバーフローが悪さをして計算結果が実行環境によって変わってしまっているものとも思われます。

[rdld036]

私のコードで114までなら分割数が正しく求められていたのですが、115以降になると違う値になりました。ネット上にあった、200までの分割数の一覧(MacMahonという人がつくったもの)をみるとそれでもまだ10^9オーダーなので、オーバーフローが関係しているのは間違いないと思うのですが、どうも納得できないところが少しあります。。。

[rdld036]

今私が使っている環境はC++17、Ubuntu 22.04.1 LTSです。

[marc2825]

200までの分割数の一覧(MacMahonという人がつくったもの)の具体的な記載がないので自分が発見した表（参照： https://fortran66.hatenablog.com/entry/2016/10/25/012637　）に基づいた解答になりますが、その表における分割数の定義が、「自然数ｎをr個の自然数（≠非負整数）に分割する組み合わせ数」となっているので、そもそも本問で考えている「分割数」とは別物（本問の方が0を含んで良い分大きくなる）ではないでしょうか？また、問題の制約も最大で1000なので仮に200時点ではオーバーフローしなくてもその後にオーバーフローすると予想できますし、自分の発見した表ではn=200におけるオーダーが10^12のオーダーになっていたので、そもそもこの時点でオーバーフローしている気がします。

[rdld036]

n = 115以降になると、P[n][n]の値が10 ^ 9を超えるからおかしくなるのですね。原因がわかりました。最後までわかりやすく教えていただきありがとうございます。